Module: Msf::DBManager::Connection

Included in:
Msf::DBManager
Defined in:
lib/msf/core/db_manager/connection.rb

Instance Method Summary collapse

Instance Method Details

#active ⇒ Object

Returns true if we are ready to load/store data



3
4
5
6
7
8
9
10
11
12
13
 
# File 'lib/msf/core/db_manager/connection.rb', line 3

def active
  # In a Rails test scenario there will be a connection established already, and it needs to be checked manually to see if a migration is required
  # This check normally happens in after_establish_connection, but that might not always get called - for instance during RSpec tests
  if Rails.env.test? && migrated.nil? && usable && connection_established?
    self.migrated = !needs_migration?
  end

  # usable and migrated a just Boolean attributes, so check those first because they don't actually contact the
  # database.
  usable && migrated && connection_established?
end

#after_establish_connection(opts = {}) ⇒ void

This method returns an undefined value.

Finishes #connect after ‘ApplicationRecord.establish_connection` has succeeded by migrating database and setting Workspace#workspace.



19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
 
# File 'lib/msf/core/db_manager/connection.rb', line 19

def after_establish_connection(opts={})
  begin
    # Migrate the database, if needed
    migrate(opts)
  rescue ::Exception => exception
    self.error = exception
    elog('DB.connect threw an exception', error: exception)

    # remove connection to prevent issues when re-establishing connection
    ApplicationRecord.remove_connection
  else
    # Flag that migration has completed
    self.migrated = true
  end
end

#connect(opts = {}) ⇒ Object

Connects this instance to a database



38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
 
# File 'lib/msf/core/db_manager/connection.rb', line 38

def connect(opts={})
  return false if not @usable

  nopts = opts.dup
  if (nopts['port'])
    nopts['port'] = nopts['port'].to_i
  end

  # Prefer the config file's pool setting
  nopts['pool'] ||= 75

  # Prefer the config file's wait_timeout setting too
  nopts['wait_timeout'] ||= 300

  begin
    # Check ApplicationRecord was already connected by Rails::Application.initialize! or some other API.
    unless connection_established?
      create_db(nopts)

      # Configure the database adapter
      ApplicationRecord.establish_connection(nopts)
    end
  rescue ::Exception => e
    self.error = e
    elog('DB.connect threw an exception', error: e)
    return false
  ensure
    after_establish_connection(nopts)
  end

  true
end

#connection_established? ⇒ true, false

Checks if the spec passed to ‘ApplicationRecord.establish_connection` can connect to the database.

Returns:

  • (true) —

    if an active connection can be made to the database using the current config.

  • (false) —

    if an active connection cannot be made to the database.



116
117
118
119
120
121
122
123
124
125
 
# File 'lib/msf/core/db_manager/connection.rb', line 116

def connection_established?
  begin
    # use with_connection so the connection doesn't stay pinned to the thread.
    ApplicationRecord.connection_pool.with_connection {
      ApplicationRecord.connection.active?
    }
  rescue ActiveRecord::ConnectionNotEstablished, PG::ConnectionBad => error
    false
  end
end

#create_db(opts) ⇒ Object

Attempt to create the database

If the database already exists this will fail and we will continue on our merry way, connecting anyway. If it doesn’t, we try to create it. If that fails, then it wasn’t meant to be and the connect will raise a useful exception so the user won’t be in the dark; no need to raise anything at all here.



81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
 
# File 'lib/msf/core/db_manager/connection.rb', line 81

def create_db(opts)
  begin
    case opts["adapter"]
    when 'postgresql'
      # Try to force a connection to be made to the database, if it succeeds
      # then we know we don't need to create it :)
      ApplicationRecord.establish_connection(opts)
      # Do the checkout, checkin dance here to make sure this thread doesn't
      # hold on to a connection we don't need
      conn = ApplicationRecord.connection_pool.checkout
      ApplicationRecord.connection_pool.checkin(conn)
    end
  rescue ::Exception => e
    errstr = e.to_s
    if errstr =~ /does not exist/i or errstr =~ /Unknown database/
      ilog("Database doesn't exist \"#{opts['database']}\", attempting to create it.")
      ApplicationRecord.establish_connection(
          opts.merge(
              'database' => 'postgres',
              'schema_search_path' => 'public'
          )
      )

      ApplicationRecord.connection.create_database(opts['database'])
    else
      ilog("Trying to continue despite failed database creation: #{e}")
    end
  end
  ApplicationRecord.remove_connection
end

#disconnect ⇒ Object

Disconnects a database session



130
131
132
133
134
135
136
137
138
139
 
# File 'lib/msf/core/db_manager/connection.rb', line 130

def disconnect
  begin
    ApplicationRecord.remove_connection
    self.migrated = nil
    self.modules_cached = false
  rescue ::Exception => e
    self.error = e
    elog('DB.disconnect threw an exception:', error: e)
  end
end