infosecn1nja
/
metasploit-framework
mirror of https://github.com/infosecn1nja/metasploit-framework.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add db_workspace command & other db refactoring. 

* Added "workspaces" table and associated ActiveRecord class.
 * Moved ActiveRecord models from db_objects.rb into separate files.
 * Do the DB migration check every time you connect (was previously done
   during db_create).
 * Use :dependent => :destroy associations so that we don't have to
   manually delete the dependent objects.

git-svn-id: file:///home/svn/framework3/trunk@7861 4d416f70-5f16-0410-b530-b9f4589650da
unstable
Mike Smith 2009-12-14 22:52:34 +00:00
parent 8317b69aca
commit f9ffc8b8bc
15 changed files with 258 additions and 119 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
data/sql/migrate/002_add_workspaces.rb Normal file
View File

@ -0,0 +1,29 @@
class AddWorkspaces < ActiveRecord::Migration
def self.up
create_table :workspaces do |t|
t.string :name
t.timestamps
end
change_table :hosts do |t|
t.integer :workspace_id, :required => true
end
remove_index :hosts, :column => :address
w = Msf::DBManager::Workspace.default
Msf::DBManager::Host.update_all ["workspace_id = ?", w.id]
end
def self.down
drop_table :workspaces
change_table :hosts do |t|
t.remove :workspace_id
end
add_index :hosts, :address, :unique => true
end
end

57
lib/msf/core/db.rb
View File

@ -207,13 +207,16 @@ class DBManager
return port
end
def workspaces
Workspace.find(:all)
end
#
# This method iterates the hosts table calling the supplied block with the
# host instance of each entry.
#
def each_host(&block)
Host.find_each do |host|
workspace.hosts.each do |host|
block.call(host)
end
end
@ -222,7 +225,7 @@ class DBManager
# This methods returns a list of all hosts in the database
#
def hosts
Host.find(:all)
workspace.hosts
end
#
@ -286,17 +289,36 @@ class DBManager
Note.find(:all)
end
def default_workspace
Workspace.default
end
def find_workspace(name)
Workspace.find_by_name(name)
end
#
# Creates a new workspace in the database
#
def add_workspace(context, name)
Workspace.find_or_create_by_name(name)
end
#
# Find or create a host matching this address/comm
#
def get_host(context, address, comm='')
if comm.length > 0
host = Host.find(:first, :conditions => [ "address = ? and comm = ?", address, comm])
host = workspace.hosts.find_by_address_and_comm(address, comm)
else
host = Host.find(:first, :conditions => [ "address = ? ", address ])
host = workspace.hosts.find_by_address(address)
end
if (not host)
host = Host.create(:address => address, :comm => comm, :state => HostState::Unknown, :created => Time.now)
host = workspace.hosts.create(
:address => address,
:comm => comm,
:state => HostState::Unknown,
:created => Time.now)
framework.events.on_db_host(context, host)
end
@ -404,19 +426,8 @@ class DBManager
# Deletes a host and associated data matching this address/comm
#
def del_host(context, address, comm='')
host = Host.find(:first, :conditions => ["address = ? and comm = ?", address, comm])
return unless host
services = Service.find(:all, :conditions => ["host_id = ?", host[:id]]).map { |s| s[:id] }
services.each do |sid|
Vuln.delete_all(["service_id = ?", sid])
Service.delete(sid)
end
Note.delete_all(["host_id = ?", host[:id]])
Host.delete(host[:id])
host = workspace.hosts.find(:first, :conditions => ["address = ? and comm = ?", address, comm])
host.destroy if host
end
#
@ -424,15 +435,9 @@ class DBManager
#
def del_service(context, address, proto, port, comm='')
host = get_host(context, address, comm)
return unless host
services = Service.find(:all, :conditions => ["host_id = ? and proto = ? and port = ?", host[:id], proto, port]).map { |s| s[:id] }
services.each do |sid|
Vuln.delete_all(["service_id = ?", sid])
Service.delete(sid)
end
host.services.find(:all, :conditions => { :proto => proto, :port => port}).destroy_all
end
#
@ -453,7 +458,7 @@ class DBManager
# Look for an address across all comms
#
def has_host?(addr)
Host.find_by_address(addr)
workspace.hosts.find_by_address(addr)
end
#

10
lib/msf/core/db_manager.rb
View File

@ -27,6 +27,9 @@ class DBManager
# Returns the active driver
attr_accessor :driver
# Returns the active workspace
attr_accessor :workspace
# Stores the error message for why the db was not loaded
attr_accessor :error
@ -55,6 +58,7 @@ class DBManager
require 'active_record'
require 'active_support'
require 'msf/core/db_objects'
require 'msf/core/model'
@usable = true
rescue ::Exception => e
@ -117,6 +121,12 @@ class DBManager
begin
# Configure the database adapter
ActiveRecord::Base.establish_connection(nopts)
# Migrate the database, if needed
migrate
# Set the default workspace
framework.db.workspace = framework.db.default_workspace
rescue ::Exception => e
self.error = e
elog("DB.connect threw an exception: #{e}")

82
lib/msf/core/db_objects.rb
View File

@ -1,5 +1,6 @@
module Msf
##
#
# This module defines all of the DB database tables
@ -39,87 +40,6 @@ module DBSave
end
# Host object definition
class Host < ActiveRecord::Base
include DBSave
has_many :services
has_many :clients
has_many :vulns
end
class Client < ActiveRecord::Base
include DBSave
belongs_to :host
def host
Host.find(:first, :conditions => [ "id = ?", host_id ])
end
end
# Service object definition
class Service < ActiveRecord::Base
include DBSave
has_many :vulns
belongs_to :host
def host
Host.find(:first, :conditions => [ "id = ?", host_id ])
end
end
# Vuln object definition
class Vuln < ActiveRecord::Base
include DBSave
belongs_to :host
belongs_to :service
has_and_belongs_to_many :refs, :join_table => :vulns_refs
def service
Service.find(:first, :conditions => [ "id = ?", service_id ])
end
def host
Host.find(:first, :conditions => [ "id = ?", host_id ])
end
end
# Reference object definition
class Ref < ActiveRecord::Base
include DBSave
has_and_belongs_to_many :vulns, :join_table => :vulns_refs
end
# Reference object definition
class VulnRefs < ActiveRecord::Base
set_table_name 'vulns_refs'
include DBSave
end
# Service object definition
class Note < ActiveRecord::Base
include DBSave
belongs_to :host
def host
Host.find(:first, :conditions => [ "id = ?", host_id ])
end
end
# WMAP Request object definition
class WmapRequest < ::ActiveRecord::Base
include DBSave
# Magic.
end
# WMAP Target object definition
class WmapTarget < ::ActiveRecord::Base
include DBSave
# Magic.
end
end
end

10
lib/msf/core/model.rb Normal file
View File

@ -0,0 +1,10 @@
require 'msf/core/model/client'
require 'msf/core/model/host'
require 'msf/core/model/note'
require 'msf/core/model/ref'
require 'msf/core/model/service'
require 'msf/core/model/workspace'
require 'msf/core/model/vuln'
require 'msf/core/model/wmap_target'
require 'msf/core/model/wmap_request'

10
lib/msf/core/model/client.rb Normal file
View File

@ -0,0 +1,10 @@
module Msf
class DBManager
class Client < ActiveRecord::Base
include DBSave
belongs_to :host
end
end
end

17
lib/msf/core/model/host.rb Normal file
View File

@ -0,0 +1,17 @@
module Msf
class DBManager
class Host < ActiveRecord::Base
include DBSave
belongs_to :workspace
has_many :services, :dependent => :destroy
has_many :clients, :dependent => :destroy
has_many :vulns, :dependent => :destroy
has_many :notes, :dependent => :destroy
validates_uniqueness_of :address, :scope => :workspace_id
end
end
end

10
lib/msf/core/model/note.rb Normal file
View File

@ -0,0 +1,10 @@
module Msf
class DBManager
class Note < ActiveRecord::Base
include DBSave
belongs_to :host
end
end
end

10
lib/msf/core/model/ref.rb Normal file
View File

@ -0,0 +1,10 @@
module Msf
class DBManager
class Ref < ActiveRecord::Base
include DBSave
has_and_belongs_to_many :vulns, :join_table => :vulns_refs
end
end
end

11
lib/msf/core/model/service.rb Normal file
View File

@ -0,0 +1,11 @@
module Msf
class DBManager
class Service < ActiveRecord::Base
include DBSave
has_many :vulns, :dependent => :destroy
belongs_to :host
end
end
end

12
lib/msf/core/model/vuln.rb Normal file
View File

@ -0,0 +1,12 @@
module Msf
class DBManager
class Vuln < ActiveRecord::Base
include DBSave
belongs_to :host
belongs_to :service
has_and_belongs_to_many :refs, :join_table => :vulns_refs
end
end
end

11
lib/msf/core/model/wmap_request.rb Normal file
View File

@ -0,0 +1,11 @@
module Msf
class DBManager
# WMAP Request object definition
class WmapRequest < ::ActiveRecord::Base
include DBSave
# Magic.
end
end
end

11
lib/msf/core/model/wmap_target.rb Normal file
View File

@ -0,0 +1,11 @@
module Msf
class DBManager
# WMAP Target object definition
class WmapTarget < ::ActiveRecord::Base
include DBSave
# Magic.
end
end
end

24
lib/msf/core/model/workspace.rb Normal file
View File

@ -0,0 +1,24 @@
module Msf
class DBManager
class Workspace < ActiveRecord::Base
include DBSave
DEFAULT = "default"
has_many :hosts, :dependent => :destroy
validates_uniqueness_of :name
validates_presence_of :name
def self.default
Workspace.find_or_create_by_name(DEFAULT)
end
def default?
name == DEFAULT
end
end
end
end

73
lib/msf/ui/console/command_dispatcher/db.rb
View File

@ -44,6 +44,7 @@ class Db
}
more = {
"db_workspace" => "Switch between database workspaces",
"db_hosts" => "List all hosts in the database",
"db_services" => "List all services in the database",
"db_vulns" => "List all vulnerabilities in the database",
@ -64,6 +65,66 @@ class Db
framework.db.active ? base.merge(more) : base
end
def cmd_db_workspace(*args)
while (arg = args.shift)
case arg
when '-h','--help'
print_line("Usage:")
print_line(" db_workspace List workspaces")
print_line(" db_workspace [name] Switch workspace")
print_line(" db_workspace -a [name] Add workspace")
print_line(" db_workspace -d [name] Delete workspace")
print_line(" db_workspace -h Show this help information")
return
when '-a','--add'
adding = true
when '-d','--del'
deleting = true
else
name = arg
end
end
if adding and name
# Add workspace
workspace = framework.db.add_workspace(nil, name)
print_status("Added workspace: #{workspace.name}")
framework.db.workspace = workspace
elsif deleting and name
# Delete workspace
workspace = framework.db.find_workspace(name)
if workspace.nil?
print_error("Workspace not found: #{name}")
elsif workspace.default?
print_error("Can't delete default workspace")
else
workspace.destroy
print_status("Deleted workspace: #{name}")
framework.db.workspace = framework.db.default_workspace if framework.db.workspace == workspace
end
elsif name
# Switch workspace
workspace = framework.db.find_workspace(name)
if workspace
framework.db.workspace = workspace
print_status("Workspace: #{workspace.name}")
else
print_error("Workspace not found: #{name}")
return
end
else
# List workspaces
framework.db.workspaces.each do |s|
pad = (s == framework.db.workspace) ? "* " : " "
print_line(pad + s.name)
end
end
end
def cmd_db_workspace_tabs(str, words)
framework.db.workspaces.map { |s| s.name } if (words & ['-a','--add']).empty?
end
def cmd_db_hosts(*args)
onlyup = false
hosts = nil
@ -1169,10 +1230,6 @@ class Db
raise RuntimeError.new("Failed to connect to the database: #{framework.db.error}")
end
if (not framework.db.migrate)
raise RuntimeError.new("Failed to create database schema: #{framework.db.error}")
end
print_status("Successfully connected to the database")
print_status("File: #{opts['dbfile']}")
@ -1310,10 +1367,6 @@ class Db
raise RuntimeError.new("Failed to connect to the database: #{framework.db.error}")
end
if (not framework.db.migrate)
raise RuntimeError.new("Failed to create database schema: #{framework.db.error}")
end
print_status("Database creation complete (check for errors)")
end
@ -1524,10 +1577,6 @@ class Db
raise RuntimeError.new("Failed to connect to the database: #{framework.db.error}")
end
if (not framework.db.migrate)
raise RuntimeError.new("Failed to create database schema: #{framework.db.error}")
end
print_status("Database creation complete (check for errors)")
end