#!/usr/bin/env ruby

################################
################################
#   WARNING: FILE DEPRECATED   #
################################
################################

## The following variables will ALWAYS be replaced:
## ${PCAP_FILE} - The pcap datafile
## ${PCAP_BASE} - The basename of the pcap datafile
## ${MODULE_DIR} - The base module directory (modules)
## ${MODULE_NAME} - The base name (without the .module) of the module
## ${OUTPUT_DIR} - The base output directory


$MODULE_DIR = "modules"
$DEBUG = false

$modules = []
$datafile = ""
$basefile = ""
$outputdir = "output"

## NSM_Module is the class encapsulating each module
class NSM_Module
  
  def initialize(name,desc,defs,cmds)
    # mod_name - String
    @mod_name = name
    # enabled - boolean
    @enabled = true
    # description - String
    @description = desc
    # defaults - Hash
    @defaults = parse_defaults(defs)
    # commands - Array
    @commands = cmds
  end
  
  def enabled?
    @enabled == true
  end
  
  def get_name
    return @mod_name
  end
  
  def get_defaults
    return @defaults
  end
  
  def get_commands
    return @commands
  end
  
  def set_enabled(en)
    @enabled = en
  end
  
  def dump
    if self.enabled?
      print "[+] "
    else
      print "[-] "
    end
    puts "#{@mod_name} - #{@description}"
  end

  def print_info
    mname = self.get_name
    filename = "#{$MODULE_DIR}/#{mname}.module/info"
    if File.exist?(filename)
      system("cat #{filename}")
    else
      puts "No info file found."
    end
  end

  def update_option(name,val)
    if name.length < 1 or val.length < 1
      puts "Invalid option or value"
      return
    end
    
    defs = self.get_defaults
    if defs[name].nil?
      puts "The variable #{name} is not settable."
      return
    end
    puts "Setting ${#{name.upcase}} = #{val}"
    defs[name] = val
  end

  def run_commands
    @commands.each { |cmd|
      dirname = $outputdir + "/" + self.get_name
      mod_name = self.get_name
      
      ## Don't replace the original command, otherwise can't change filename
      new_cmd = cmd
      ## Do the most important 3 first
      new_cmd.gsub!(/\$\{PCAP_FILE\}/i,$datafile)
      new_cmd.gsub!(/\$\{PCAP_BASE\}/i,$basefile)
      new_cmd.gsub!(/\$\{MODULE_DIR\}/i,$MODULE_DIR)
      new_cmd.gsub!(/\$\{MODULE_NAME\}/i,mod_name)
      new_cmd.gsub!(/\$\{OUTPUT_DIR\}/i,dirname)
      
      ## Replace others with options
      defs = self.get_defaults
      defs.each_pair { |d, val|
        ## This allows the use of these 3 vars in the defaults file
        new_val = val
        new_val.gsub!(/\$\{PCAP_FILE\}/i,$datafile)
        new_val.gsub!(/\$\{PCAP_BASE\}/i,$basefile)
        new_val.gsub!(/\$\{MODULE_DIR\}/i,$MODULE_DIR)
        new_val.gsub!(/\$\{MODULE_NAME\}/i,mod_name)
        new_val.gsub!(/\$\{OUTPUT_DIR\}/i,dirname)
        
        new_cmd.gsub!(/\$\{#{d}\}/i,new_val)
        
      }
      
      ## Build the output directories if they don't exist
      if File.exist?(dirname)
        if File.directory?(dirname)
          puts "Directory #{dirname} already exists, not recreating"
        else
          puts "Error: #{dirname} exists, but is not a directory."
          return
        end
      else
        puts "Creating directory #{dirname}"
        res = system("mkdir -p #{dirname}")
        if res
          puts "-> Success."
        else
          puts "-> Failure."
        end
      end
      
      ## Change this to actually execute command
      puts "--> #{new_cmd.to_s}"
      res = system("#{new_cmd.to_s}")
      if res
        puts "-> Success."
      else
        puts "-> Failure."
      end
    }
  end

  private
  def parse_defaults(defs)
    new_defs = {}
    defs.each { |d|
      a = d.split(/=/)
      new_defs.store(a[0],a[1])
    }
    return new_defs
  end

end

## Start of helping methods
## Print out all modules loaded
def list_all_modules
  $modules.each { |mod|
    mod.dump
  }
  puts "--------------------------"
  puts "+ = Enabled\n- = Disabled"
end

## Return the description of the <name> module
def get_description(name)
  if File.exist?("#{$MODULE_DIR}/#{name}/description")
    desc = `cat #{$MODULE_DIR}/#{name}/description`
  else
    desc = "No description"
  end
  return desc.chomp
end

## Return the default values for the <name> module
def get_defaults(name)
  defs = []
  if File.exist?("#{$MODULE_DIR}/#{name}/defaults")
    begin
      file = File.open("#{$MODULE_DIR}/#{name}/defaults", "r")
      file.each_line {
        |line|
        defs.push(line.chomp)
      }
      file.close
    rescue
      puts "\nUnable to read file #{$MODULE_DIR}/#{name}/defaults: #{$!}\n"
    end
  end
  return defs
end

## Return the commands for the <name> module
def get_commands(name)
  cmds = []
  if File.exist?("#{$MODULE_DIR}/#{name}/#{name.gsub(/\.module$/,"")}")
    begin
      file = File.open("#{$MODULE_DIR}/#{name}/#{name.gsub(/\.module$/,"")}", "r")
      file.each_line {
        |line|
        cmds.push(line.chomp)
      }
      file.close
    rescue
      puts "\nUnable to read file #{$MODULE_DIR}/#{name}/#{name.gsub(/\.module$/,"")}: #{$!}\n"
    end
  end
  return cmds
end

## Return the module given by <mod_name>
def get_mod_by_name(mod_name)
  $modules.each { |mod|
    if mod.get_name == mod_name
      return mod
    end
  }
  return nil
end

## Print the lobster banner :D
def print_banner
  ## New header :D
  puts "=-" * 30
  puts "                         ,.---."   
  puts "               ,,,,     /    _ `."
  puts "                \\\\\\\\   /      \\  )"
  puts "                 |||| /\\/``-.__\\/          NSM"
  puts "                 ::::/\\/_"
  puts " {{`-.__.-'(`(^^(^^^(^ 9 `.========='"
  puts "{{{{{{ { ( ( (  (   (-----:="
  puts " {{.-'~~'-.(,(,,(,,,(__6_.'=========."
  puts "                 ::::\\/\\ "
  puts "                 |||| \\/\\  ,-'/\\           console"
  puts "                ////   \\ `` _/  )"
  puts "               ''''     \\  `   /"
  puts "                         `---''"
  puts "=-" * 30
  print "\n"
end

## Load the modules in <dir>
def load_modules(dir)
  $modules = []

  if !File.directory?(dir)
    puts "Error: Module directory '#{dir}' is not a directory or does not exist."
    puts "No modules will be loaded."
    return
  end
  
  mod_num = 0
  mods = Dir.new(dir)
  ## Load modules
  mods.each { |mod_name|
    if mod_name =~ /\w\.module$/i
      print "Loading #{mod_name}..."
      m = NSM_Module.new(mod_name.gsub(/\.module$/,""),
                         get_description(mod_name),
                         get_defaults(mod_name),
                         get_commands(mod_name))
      $modules.push(m)
      puts "done."
      mod_num += 1
    end
  }
  puts "#{mod_num} modules loaded.\n\n"
end

##################################
###### Command declarations ######
##################################
## Command declarations are functions that get called
## By the NSM shell, hence <function>_cmd notation.
## They *always* have 1 argument, which is a string
## containing all the arguments from the nsm shell.

def help_cmd(args)
  puts "Available Commands:\n\n"
  puts " file, help, info, list, options, output, run, set, toggle, quit"
  print "\n"
  puts " Command\t\t  Description"
  puts "-" * 70
  puts " file <filename>\t- Set the ${PCAP_FILE} filename (the file to be analyzed)."
  puts " help\t\t\t- Display this help."
  puts " info <module>\t\t- Display information about the <module> module."
  puts " list\t\t\t- List modules and their toggled status [+] is enabled, [-] is disabled."
  puts " options [<module]\t- Display global options. If <module> is given, display that"
  puts " \t\t\t  module's options as well (ex: 'options hash'). 'options' will also display"
  puts " \t\t\t  the commands that will be run (before substitution) for that module."
  puts " output <dir>\t\t- Set the output directory (default: 'output')."
  puts " run\t\t\t- Run the enabled analysis modules."
  puts " set <mod> <opt> <val>\t- Set the <mod> module's option '<opt>' to have the value <val>. "
  puts " \t\t\t  (ex: 'set aimsnarf OUTPUT_FILE ${PCAP_BASE}.aim.txt')."
  puts " toggle <module>\t- Toggle <module> to be enabled or disabled. Default is enabled."
  puts " \t\t\t  Use 'toggle none' or 'toggle all' to turn enable/disable all modules."
  puts " quit\t\t\t- Quit the NSM console."
  
  print "\n"
  return
end

def list_cmd(args)
  list_all_modules
end

def toggle_cmd(name)
  if name.length < 1
    puts "Need a module name"
    return
  end
  if name == "all"
    $modules.each { |mod| mod.set_enabled(true); }
    puts "All modules turned on."
    return
  elsif name == "none"
    $modules.each { |mod| mod.set_enabled(false); }
    puts "All modules turned off."
    return
  else
    $modules.each { |mod|
      if mod.get_name() =~ /#{name}/
        if mod.enabled?
          mod.set_enabled(false)
          puts "#{name} module turned off."
        else
          mod.set_enabled(true)
          puts "#{name} module turned on."
        end
        return
      end
    }
  end
  puts "Module '#{name}' not found."
  return
end

def options_cmd(name)
  puts "\nGlobal options:"
  puts "-" * 35
  puts "${PCAP_FILE}: #{$datafile}"
  puts "${PCAP_BASE}: #{$basefile}"
  puts "${OUTPUT_DIR}: #{$outputdir}"
  puts "${MODULE_DIR}: #{$MODULE_DIR}\n\n"
    
  return if name.length < 1

  mod = get_mod_by_name(name)
  if mod.nil?
    puts "Unknown module #{name}"
    return
  end
  defs = mod.get_defaults
  puts "Options for module #{name}:"
  puts "-" * 35
  defs.each_pair { |d, val|
    puts "${#{d}} = #{val}"
  }
  cmds = mod.get_commands
  puts "\nCommand(s) to be executed for module #{name}:"
  puts "-" * 35
  cmds.each { |c| puts c }
  print "\n"
end

def file_cmd(name)
  if name.length < 1
    puts "Need a pcap filename"
    return
  end
  $datafile = name
  puts "Setting ${PCAP_FILE} = #{$datafile}"
  $basefile = `basename #{name}`.chomp
  puts "Setting ${PCAP_BASE} = #{$basefile}"
end

def output_cmd(name)
  if name.length < 1
    puts "Need a output directory name"
    return
  end
  puts "Setting ${OUTPUT_DIR} = #{name}"
  $outputdir = name
end

def run_cmd(args)
      
  if $datafile == ""
    puts "Error: no pcap datafile has been specified"
    puts "specify one with 'file <file.pcap>'"
    return
  end
  
  if !File.exist?($datafile)
    puts "ERROR: #{$datafile} does not exist. Please set file using 'file <filename>'."
    return
  end
  
  puts "\nExecuting analysis...\n\n"
  
  $modules.each { |mod|
    if mod.enabled?
      puts "===> module #{mod.get_name} running..."
      mod.run_commands 
      puts "===> module #{mod.get_name} finished."
    else
      puts "===> module #{mod.get_name} skipped."
    end
  }
  print "\n"
end

def info_cmd(name)
  mod = get_mod_by_name(name)
  if mod.nil?
    puts "Module not found."
    return
  end
  mod.print_info
end

def set_cmd(args)
  if args.length < 1
    puts "Need a module name."
    return
  end
  arglist = args.split(/ /)
  mod = get_mod_by_name(arglist[0])
  if mod.nil?
    puts "Module #{arglist[0]} not found."
    return
  end
  if arglist[1].nil?
    puts "Need variable name (ex: OUTPUT_DIR)"
    return
  end
  if arglist[2].nil?
    puts "Need variable value (ex: file.pcap)"
    return
  end
  mod.update_option(arglist[1],arglist[2])
end

##################################
######## End declarations ########
##################################

###################
## Begin program ##
###################

## Print the lobster
print_banner()

## Set {$PCAP_FILE} if passed in as an argument
file_cmd(ARGV[0]) if ARGV.length > 0

## Load modules
load_modules($MODULE_DIR)

puts "Default ${OUTPUT_DIR} is '#{$outputdir}'"

cmd = ""

print "\n"
puts "=-" * 35
puts "Welcome to NSM Console, type 'help' to see available commands"
puts "=-" * 35
puts "Note: All modules are enabled by default, use 'list' to list available"
puts "modules and 'toggle <module>' to disable/enable a module.\n\n"

## Command loop until quit
while (cmd !~ /^(quit|q_cmd)/)
  print "nsm> "
  cmd = $stdin.gets

  ## Since we use eval (for now), ';' is very very bad.
  if cmd =~ /;/
    puts "Bad! ';' can't be used! (yet)"
    cmd = ""
  end

  ## Split the command into function and arguments
  cmd.gsub!(/(\w)\s(.*)/) { $1 }
  cmd.chomp!
  args = $2
  
  cmd += "_cmd(\"#{args.to_s}\")"
  begin
    eval(cmd) unless cmd =~ /^(quit_cmd|q_cmd|_cmd)/
  rescue
    puts "Error evaluating command (Bad command?)"
    puts "Error: #{$!}" if $DEBUG
  end

end
