From f9ffc8b8bcf78003ea803bbe0dc563eaeee46c70 Mon Sep 17 00:00:00 2001 From: Mike Smith Date: Mon, 14 Dec 2009 22:52:34 +0000 Subject: [PATCH] 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 --- data/sql/migrate/002_add_workspaces.rb | 29 ++++++++ lib/msf/core/db.rb | 57 +++++++------- lib/msf/core/db_manager.rb | 10 +++ lib/msf/core/db_objects.rb | 82 +-------------------- lib/msf/core/model.rb | 10 +++ lib/msf/core/model/client.rb | 10 +++ lib/msf/core/model/host.rb | 17 +++++ lib/msf/core/model/note.rb | 10 +++ lib/msf/core/model/ref.rb | 10 +++ lib/msf/core/model/service.rb | 11 +++ lib/msf/core/model/vuln.rb | 12 +++ lib/msf/core/model/wmap_request.rb | 11 +++ lib/msf/core/model/wmap_target.rb | 11 +++ lib/msf/core/model/workspace.rb | 24 ++++++ lib/msf/ui/console/command_dispatcher/db.rb | 73 +++++++++++++++--- 15 files changed, 258 insertions(+), 119 deletions(-) create mode 100644 data/sql/migrate/002_add_workspaces.rb create mode 100644 lib/msf/core/model.rb create mode 100644 lib/msf/core/model/client.rb create mode 100644 lib/msf/core/model/host.rb create mode 100644 lib/msf/core/model/note.rb create mode 100644 lib/msf/core/model/ref.rb create mode 100644 lib/msf/core/model/service.rb create mode 100644 lib/msf/core/model/vuln.rb create mode 100644 lib/msf/core/model/wmap_request.rb create mode 100644 lib/msf/core/model/wmap_target.rb create mode 100644 lib/msf/core/model/workspace.rb diff --git a/data/sql/migrate/002_add_workspaces.rb b/data/sql/migrate/002_add_workspaces.rb new file mode 100644 index 0000000000..abb6daebb0 --- /dev/null +++ b/data/sql/migrate/002_add_workspaces.rb @@ -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 \ No newline at end of file diff --git a/lib/msf/core/db.rb b/lib/msf/core/db.rb index a7cdcad813..99aaafc104 100644 --- a/lib/msf/core/db.rb +++ b/lib/msf/core/db.rb @@ -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 # diff --git a/lib/msf/core/db_manager.rb b/lib/msf/core/db_manager.rb index b37d342a70..5c581082f0 100644 --- a/lib/msf/core/db_manager.rb +++ b/lib/msf/core/db_manager.rb @@ -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}") diff --git a/lib/msf/core/db_objects.rb b/lib/msf/core/db_objects.rb index e41199f9b8..93c0baa53e 100644 --- a/lib/msf/core/db_objects.rb +++ b/lib/msf/core/db_objects.rb @@ -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 diff --git a/lib/msf/core/model.rb b/lib/msf/core/model.rb new file mode 100644 index 0000000000..13b804d95d --- /dev/null +++ b/lib/msf/core/model.rb @@ -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' \ No newline at end of file diff --git a/lib/msf/core/model/client.rb b/lib/msf/core/model/client.rb new file mode 100644 index 0000000000..43a49f42ed --- /dev/null +++ b/lib/msf/core/model/client.rb @@ -0,0 +1,10 @@ +module Msf +class DBManager + +class Client < ActiveRecord::Base + include DBSave + belongs_to :host +end + +end +end \ No newline at end of file diff --git a/lib/msf/core/model/host.rb b/lib/msf/core/model/host.rb new file mode 100644 index 0000000000..2763d20467 --- /dev/null +++ b/lib/msf/core/model/host.rb @@ -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 \ No newline at end of file diff --git a/lib/msf/core/model/note.rb b/lib/msf/core/model/note.rb new file mode 100644 index 0000000000..eed8188dee --- /dev/null +++ b/lib/msf/core/model/note.rb @@ -0,0 +1,10 @@ +module Msf +class DBManager + +class Note < ActiveRecord::Base + include DBSave + belongs_to :host +end + +end +end \ No newline at end of file diff --git a/lib/msf/core/model/ref.rb b/lib/msf/core/model/ref.rb new file mode 100644 index 0000000000..f9eb612fcd --- /dev/null +++ b/lib/msf/core/model/ref.rb @@ -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 \ No newline at end of file diff --git a/lib/msf/core/model/service.rb b/lib/msf/core/model/service.rb new file mode 100644 index 0000000000..de07688bc6 --- /dev/null +++ b/lib/msf/core/model/service.rb @@ -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 \ No newline at end of file diff --git a/lib/msf/core/model/vuln.rb b/lib/msf/core/model/vuln.rb new file mode 100644 index 0000000000..ea6e122576 --- /dev/null +++ b/lib/msf/core/model/vuln.rb @@ -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 \ No newline at end of file diff --git a/lib/msf/core/model/wmap_request.rb b/lib/msf/core/model/wmap_request.rb new file mode 100644 index 0000000000..96ba3f6067 --- /dev/null +++ b/lib/msf/core/model/wmap_request.rb @@ -0,0 +1,11 @@ +module Msf +class DBManager + +# WMAP Request object definition +class WmapRequest < ::ActiveRecord::Base + include DBSave + # Magic. +end + +end +end \ No newline at end of file diff --git a/lib/msf/core/model/wmap_target.rb b/lib/msf/core/model/wmap_target.rb new file mode 100644 index 0000000000..e7a7bbfaaa --- /dev/null +++ b/lib/msf/core/model/wmap_target.rb @@ -0,0 +1,11 @@ +module Msf +class DBManager + +# WMAP Target object definition +class WmapTarget < ::ActiveRecord::Base + include DBSave + # Magic. +end + +end +end \ No newline at end of file diff --git a/lib/msf/core/model/workspace.rb b/lib/msf/core/model/workspace.rb new file mode 100644 index 0000000000..6aa6dcf242 --- /dev/null +++ b/lib/msf/core/model/workspace.rb @@ -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 \ No newline at end of file diff --git a/lib/msf/ui/console/command_dispatcher/db.rb b/lib/msf/ui/console/command_dispatcher/db.rb index bd5e9e9f62..2e06c85455 100644 --- a/lib/msf/ui/console/command_dispatcher/db.rb +++ b/lib/msf/ui/console/command_dispatcher/db.rb @@ -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