Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place.

1
# -*- coding: binary -*-
2
# This is a compatibility layer for using the pure Ruby postgres-pr instead of
3
# the C interface of postgres.
4

5
require 'postgres_msf'
6
require 'postgres/postgres-pr/connection'
7

8
# Namespace for Metasploit branch.
9
module Msf
10
module Db
11

12
class PGconn
13
  class << self
14
    alias connect new
15
  end
16

17
  def initialize(host, port, options, tty, database, user, auth)
18
    uri =
19
    if host.nil?
20
      nil
21
    elsif host[0] != ?/
22
      "tcp://#{ host }:#{ port }"
23
    else
24
      "unix:#{ host }/.s.PGSQL.#{ port }"
25
    end
26
    @host = host
27
    @db = database
28
    @user = user
29
    @conn = PostgresPR::Connection.new(database, user, auth, uri)
30
  end
31

32
  def close
33
    @conn.close
34
  end
35

36
  attr_reader :host, :db, :user
37

38
  def query(sql)
39
    PGresult.new(@conn.query(sql))
40
  end
41

42
  alias exec query
43

44
  def transaction_status
45
    @conn.transaction_status
46
  end
47

48
  def self.escape(str)
49
    str.gsub("'","''").gsub("\\", "\\\\\\\\")
50
  end
51

52
  def notice_processor
53
    @conn.notice_processor
54
  end
55

56
  def notice_processor=(np)
57
    @conn.notice_processor = np
58
  end
59

60
  def self.quote_ident(name)
61
    %("#{name}")
62
  end
63

64
end
65

66
class PGresult
67
  include Enumerable
68

69
  EMPTY_QUERY = 0
70
  COMMAND_OK = 1
71
  TUPLES_OK = 2
72
  COPY_OUT = 3
73
  COPY_IN = 4
74
  BAD_RESPONSE = 5
75
  NONFATAL_ERROR = 6
76
  FATAL_ERROR = 7
77

78
  def each(&block)
79
    @result.each(&block)
80
  end
81

82
  def [](index)
83
    @result[index]
84
  end
85
 
86
  def initialize(res)
87
    @res = res
88
    @fields = @res.fields.map {|f| f.name}
89
    @result = @res.rows
90
  end
91

92
  # TODO: status, getlength, cmdstatus
93

94
  attr_reader :result, :fields
95

96
  def num_tuples
97
    @result.size
98
  end
99

100
  def num_fields
101
    @fields.size
102
  end
103

104
  def fieldname(index)
105
    @fields[index]
106
  end
107

108
  def fieldnum(name)
109
    @fields.index(name)
110
  end
111

112
  def type(index)
113
    # TODO: correct?
114
    @res.fields[index].type_oid
115
  end
116

117
  def size(index)
118
    raise
119
    # TODO: correct?
120
    @res.fields[index].typlen
121
  end
122

123
  def getvalue(tup_num, field_num)
124
    @result[tup_num][field_num]
125
  end
126

127
  def status
128
    if num_tuples > 0
129
      TUPLES_OK
130
    else
131
      COMMAND_OK
132
    end
133
  end
134

135
  def cmdstatus
136
    @res.cmd_tag || ''
137
  end
138

139
  # free the result set
140
  def clear
141
    @res = @fields = @result = nil
142
  end
143

144
  # Returns the number of rows affected by the SQL command
145
  def cmdtuples
146
    case @res.cmd_tag
147
    when nil 
148
      return nil
149
    when /^INSERT\s+(\d+)\s+(\d+)$/, /^(DELETE|UPDATE|MOVE|FETCH)\s+(\d+)$/
150
      $2.to_i
151
    else
152
      nil
153
    end
154
  end
155

156
end
157

158
class PGError < ::Exception
159
end
160

161
end
162
end
163

164