/*
 *  Paglo Crawler
 *  Copyright (C) 2006-2008 Paglo Labs Inc. All rights reserved.
 *  www.paglo.com
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of version 2 of the GNU General Public License as
 *  published by the Free Software Foundation.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
 */
//---------------------------------------------------------------------------
#include "RubyThread.h"
#include "AnalysisManager.h"
#include "EvidenceDescription.h"
#include "DirectoryList.h"
#include "ScanThread.h"
#include "Probes.h"
#include "CommunicationThread.h"         
#include "Debug.h"
#include "Plugin.h"
#include <ruby.h>
//---------------------------------------------------------------------------

/*
 * Global Ruby constants.
 */
VALUE PluginManager;
VALUE PluginInstaller;
VALUE NetworkDeviceClass;
VALUE PluginException;

//---------------------------------------------------------------------------

static void HandleRubyError(void)
{
  /*
   * An error occurred. Read the error message from Ruby.
   */
  VALUE lasterr = rb_gv_get("$!");
  VALUE message = rb_obj_as_string(lasterr);

  string Error = RSTRING(message)->ptr;

  if (!NIL_P(ruby_errinfo)) {
    VALUE ary = rb_funcall(ruby_errinfo, rb_intern("backtrace"), 0);
    for (int c=0; c<RARRAY(ary)->len; c++) {
      Error += string("\n    from ") + RSTRING(RARRAY(ary)->ptr[c])->ptr;
    }
  }
  LogError("RUBY ERROR: %s", Error.c_str());
  LogIssue(itPlugin, isError, "RUBY ERROR: %s", RSTRING(message)->ptr);
}

//---------------------------------------------------------------------------

TRubyThread::TRubyThread() : TPThread()
{
  InputQueue = new TEventQueue();
  CurrentPluginName = "Idle";
  PluginStartTime = time(NULL);
  Active = false;
}
//---------------------------------------------------------------------------

TRubyThread::~TRubyThread()
{
  TPEvent *Evt;
  TScanRequest *ScanReq;

  ruby_finalize();
  
  while (!InputQueue->IsEmpty()) {
    Evt = InputQueue->Get();
    if (Evt) {
      ScanReq = Evt->GetScanReq();
      delete Evt;
      if (ScanReq) {
        delete ScanReq;
      }
    }
  }
  delete InputQueue;
}
//---------------------------------------------------------------------------

void TRubyThread::Run()
{
  TPEvent *InEvt;
  TScanRequest *ScanReq;
  TScanResult *Reply;
  char *cpath;
  string str_path = _GlobalConfig.GetPluginDir();

  /*
   * Initialize the ruby interpreter and create the plugin manager.
   */
  DEBUG_MESSAGE(DEBUG_PLUGINS, ("Plugin path is: '%s'", cpath));
  if (!InitInterpreter(str_path.c_str())) {
    return;
  }

  /*
   * Get list of plugins, their descriptions, and parameters from the
   * plugin manager.
   */
  if (!GetPlugins()) {
    return;
  }

  while (!IsTerminate()) {
    /*
     * Wait for an event.
     */
    while (InputQueue->IsEmpty() && !IsTerminate()) {
      /*
       * Poll the periodic functions.
       */
      PluginManagerPoll();
    }
    PluginManagerPoll();
    
    InEvt = InputQueue->Get();

    if (InEvt) {
      if ((ScanReq = InEvt->GetScanReq()) != NULL) {
        Reply = ScanReq->PerformScan(this);
        rb_gc_start();
        TEventQueue *ReplyQueue = ScanReq->GetReplyQueue();
        if (ReplyQueue) {
          if (Reply) {
            TPEvent *OutEvt = new TPEvent(Reply);
            ReplyQueue->Put(OutEvt);
          }
        } else {
          delete Reply;
        }
        delete ScanReq;
      }
      delete InEvt;
    }
  }
}
//---------------------------------------------------------------------------

static VALUE ScannerLog(VALUE Self, VALUE Message)
{
  if (TYPE(Message) != T_NIL) {
    char *M = StringValueCStr(Message);
    if (M) {
      LogMesg("RUBY: %s", M);
    }
  }
  
  return Qnil;
}
//---------------------------------------------------------------------------

static VALUE RubyLogIssue(VALUE Self, VALUE rb_Message)
{
  if (TYPE(rb_Message) == T_STRING) {
    char *Message = StringValueCStr(rb_Message);
    if (Message) {
      LogIssue(itPlugin, isWarning, "%s", Message);
    }
  }

  return Qnil;
}
//---------------------------------------------------------------------------

static VALUE SetCurrentPlugin(VALUE Self, VALUE rb_Name)
{
  if (TYPE(rb_Name) == T_STRING) {
    char *Name = StringValueCStr(rb_Name);
    if (Name && _RubyThread->CurrentPluginName != Name) {
      _RubyThread->CurrentPluginName = Name;
      _RubyThread->PluginStartTime = time(NULL);
      _RubyThread->Active = true;
    }
  } else {
    _RubyThread->CurrentPluginName = "Idle";
    _RubyThread->Active = false;
  }

  return Qnil;
}
//---------------------------------------------------------------------------

static VALUE SubmitTree(VALUE Self, VALUE rb_Tree, VALUE rb_Structure, VALUE rb_Values)
{
  if (TYPE(rb_Tree) != T_NIL &&
      TYPE(rb_Structure) != T_NIL &&
      TYPE(rb_Values) != T_NIL) {
    char *Tree = StringValueCStr(rb_Tree);
    char *Structure = StringValueCStr(rb_Structure);
    char *Values = StringValueCStr(rb_Values);

    TSubmitTreeCmd *Cmd = new TSubmitTreeCmd(Tree, Structure, Values);
    _CommThread->EnqueueCmd(Cmd);
  } else {
    rb_raise(PluginException, "Tree submission does not include all of the required information");
  }
  
  return Qnil;
}
//---------------------------------------------------------------------------

static VALUE GetConfig(VALUE Self, VALUE rb_Param)
{
  string Param = StringValueCStr(rb_Param);
  list<TConfigParam> ParamList = _GlobalConfig.GetConfigParam(Param);
  int ParamCount = 0;
  VALUE rb_Params = rb_ary_new();
  list<TConfigParam>::iterator CurrParam;
  for (CurrParam = ParamList.begin(); CurrParam != ParamList.end(); CurrParam++) {
    string S;
    for (int i = 0; i < (CurrParam->NumFields() - 1); i++) {
      S += CurrParam->Field(i) + string(",");
    }
    S += CurrParam->Field(CurrParam->NumFields() - 1);

    rb_ary_push(rb_Params, rb_str_new2(S.c_str()));
    ParamCount++;
  }


  if (ParamCount > 0) {
    return rb_Params;
  } else {
    return Qnil;
  }
}
//---------------------------------------------------------------------------

static VALUE PortOpen(VALUE Self, VALUE rb_IPAddress, VALUE rb_Proto, VALUE rb_Port)
{
  /*
   * Check types.
   */
  if (TYPE(rb_IPAddress) != T_STRING ||
      TYPE(rb_Proto) != T_SYMBOL ||
      TYPE(rb_Port) != T_FIXNUM) {
    return Qnil;
  }
  
  string IPAddrStr = StringValueCStr(rb_IPAddress);
  ID rb_ProtoID = rb_to_id(rb_Proto);
  ID rb_TcpID = rb_intern("tcp");
  ID rb_UdpID = rb_intern("udp");
  int Port = NUM2INT(rb_Port);

    
  TIPAddress IPAddr(IPAddrStr);
  TNetworkDevice *NetDev = _AnalysisMgr->GetDeviceManager()->LookupDevice(IPAddr);
  if (NetDev) {
    if (rb_ProtoID == rb_TcpID) {
      return (NetDev->HasTcpPort(Port) ? rb_eval_string("true"): rb_eval_string("false"));
    } else if (rb_ProtoID == rb_UdpID) {
      return (NetDev->HasUdpPort(Port) ? rb_eval_string("true"): rb_eval_string("false"));
    }
  }

  return Qnil;
}
//---------------------------------------------------------------------------

static VALUE NewDevice(VALUE Self, VALUE rb_IPAddress, VALUE rb_MacAddress)
{
  TNetworkDevice *NetDev = NULL;
  if (TYPE(rb_IPAddress) == T_STRING) {
    TIPAddress IPAddr(StringValueCStr(rb_IPAddress));
    if (TYPE(rb_MacAddress) == T_STRING) {
      TMacAddress MacAddr(StringValueCStr(rb_MacAddress));
      NetDev = _AnalysisMgr->LookupOrCreateDevice(IPAddr, MacAddr);
    } else {
      NetDev = _AnalysisMgr->LookupOrCreateDevice(IPAddr);
    }
  } else if (TYPE(rb_MacAddress) == T_STRING) {
    TMacAddress MacAddr(StringValueCStr(rb_MacAddress));
    NetDev = _AnalysisMgr->LookupOrCreateDevice(MacAddr);
  }

  /*
   * Instantiate and return a network network device.
   */
  return rb_funcall(NetworkDeviceClass, rb_intern("new"), 1, INT2NUM((long)NetDev));
}
//---------------------------------------------------------------------------

static VALUE GetDeviceByIP(VALUE Self, VALUE rb_IPAddress)
{
  TIPAddress IPAddr(StringValueCStr(rb_IPAddress));
  TNetworkDevice *NetDev = _AnalysisMgr->GetDeviceManager()->LookupDevice(IPAddr);

  if (NetDev) {
    return rb_funcall(NetworkDeviceClass, rb_intern("new"), 1, INT2NUM((long)NetDev));
  } else {
    return Qnil;
  }
}
//---------------------------------------------------------------------------

static VALUE GetGUID(VALUE Self)
{
  string GUID = _GlobalConfig.GetGUID();
  return rb_str_new2(GUID.c_str());
}

//---------------------------------------------------------------------------

static VALUE GetLogFilename(VALUE Self)
{
  string log_file_name = _GlobalConfig.GetLogFilename();
  return rb_str_new2(log_file_name.c_str());
}

//---------------------------------------------------------------------------

static VALUE GetPluginDir(VALUE Self)
{
  string plugin_dir = _GlobalConfig.GetPluginDir();
  return rb_str_new2(plugin_dir.c_str());
}
//---------------------------------------------------------------------------

static VALUE FindEvidence(VALUE Self, VALUE NetDevHandle, VALUE Characteristic)
{
  TNetworkDevice *NetDev = (TNetworkDevice*)NUM2INT(NetDevHandle);

  /*
   * Check that the NetDev is valid. This is slow since the device manager is
   * doing a linear search.
   */
  TDeviceManager *DevMgr = _AnalysisMgr->GetDeviceManager();
  if (!DevMgr->LookupDevice(NetDev)) {
    rb_raise(PluginException, "Network device handle not valid in find_evidence()");
    return Qnil;
  }

  list<TMultiTypeMapValue> Evidence = NetDev->FindEvidence(STR2CSTR(Characteristic));
  VALUE ReturnList = rb_ary_new();
  for (list<TMultiTypeMapValue>::iterator Iter = Evidence.begin();
    Iter != Evidence.end(); Iter++) {
    TMultiTypeMapValue Val = *Iter;

    rb_ary_push(ReturnList, rb_str_new2(Val.Print().c_str()));
  }

  return ReturnList;
}
//---------------------------------------------------------------------------

static VALUE ScanDevice(VALUE Self, VALUE rb_IPAddress)
{
  try {
    TIPAddress IPAddr(StringValueCStr(rb_IPAddress));
    TDeviceManager *DevMgr = _AnalysisMgr->GetDeviceManager();
    TNetworkDevice *NetDev = DevMgr->LookupDevice(IPAddr);

    if (NetDev) {
      /*
       * See if a scan is in progress.
       */
      if (NetDev->NumPendingScans() == 0) {
        _AnalysisMgr->IssueScans(NetDev, false);
      }
    }
  } catch (...) {}

  return Qnil;
}
//---------------------------------------------------------------------------

static VALUE GetLocalIP(VALUE Self)
{
  string LocalIP = _GlobalConfig.GetInterfaceIPAddress().Print();
  return rb_str_new2(LocalIP.c_str());
}
//---------------------------------------------------------------------------

static VALUE ProtectedPluginManagerPoll(VALUE Arg)
{
  rb_funcall(PluginManager, rb_intern("periodic_poll"), 0);
  return Qnil;
}
//---------------------------------------------------------------------------

void TRubyThread::PluginManagerPoll()
{
  int Error;

  rb_protect(ProtectedPluginManagerPoll, 0, &Error);
  if (Error) {
    HandleRubyError();
  }
}
//---------------------------------------------------------------------------

static VALUE ProtectedInitInterpreter(VALUE RubyPath)
{
  char *Path = STR2CSTR(RubyPath);
  char include1[256], include2[256], include3[256], include4[256];
  char include5[256];
  
  sprintf(include1, "-I%s%s", Path, "Plugins");
  sprintf(include2, "-I%s%s/include", Path, "Plugins");
  sprintf(include3, "-I%s%s/include/lib", Path, "Plugins");
  sprintf(include4, "-I%s%s/include/lib/i386-mswin32", Path, "Plugins");
  sprintf(include5, "-I%s%s/lib", Path, "Plugins");
  
  /* TODO: Add the linux library path here if necessary. */

  char *argv[6];
  argv[0] = 0;
  argv[1] = include1;
  argv[2] = include2;
  argv[3] = include3;
  argv[4] = include4;    
  argv[5] = include5;

  ruby_options(6, argv);

  /*
   * Define the C++ callback functions.
   */
  rb_define_global_function("scanner_log", (VALUE(*)(ANYARGS))ScannerLog, 1);
  rb_define_global_function("log_issue", (VALUE(*)(ANYARGS))RubyLogIssue, 1);
  rb_define_global_function("submit_tree", (VALUE(*)(ANYARGS))SubmitTree, 3);
  rb_define_global_function("get_config", (VALUE(*)(ANYARGS))GetConfig, 1);
  rb_define_global_function("port_open?", (VALUE(*)(ANYARGS))PortOpen, 3);
  rb_define_global_function("new_device", (VALUE(*)(ANYARGS))NewDevice, 2);
  rb_define_global_function("get_guid", (VALUE(*)(ANYARGS))GetGUID, 0);
  rb_define_global_function("get_device_by_ip", (VALUE(*)(ANYARGS))GetDeviceByIP, 1);
  rb_define_global_function("find_evidence", (VALUE(*)(ANYARGS))FindEvidence, 2);
  rb_define_global_function("scan_device", (VALUE(*)(ANYARGS))ScanDevice, 1);
  rb_define_global_function("get_local_ip", (VALUE(*)(ANYARGS))GetLocalIP, 0);
  rb_define_global_function("set_current_plugin", (VALUE(*)(ANYARGS))SetCurrentPlugin, 1);
  rb_define_global_function("get_log_file_name", (VALUE(*)(ANYARGS))GetLogFilename, 0);
  rb_define_global_function("get_plugin_dir", (VALUE(*)(ANYARGS))GetPluginDir, 0);
  
  /*
   * Create the plugin manager.
   */
  if (!rb_require("plugin_manager")) {
    rb_raise(rb_intern("Exception"), "Unable to load the plugin manager");
  }
  VALUE PagloModule = rb_const_get(rb_cModule, rb_intern("Paglo"));
  PluginException = rb_define_class_under(PagloModule, "PluginException", rb_eStandardError);
  VALUE PluginManagerClass = rb_const_get(PagloModule, rb_intern("PluginManager"));
  VALUE PluginManager = rb_funcall(PluginManagerClass, rb_intern("new"), 0, 0);
  NetworkDeviceClass = rb_const_get(PagloModule, rb_intern("NetworkDevice"));

  /*
   * Create the plugin installer.
   */
  if (!rb_require("plugin_installer")) {
    rb_raise(rb_intern("Exception"), "Unable to load the plugin installer");
  }
  VALUE PluginInstallerClass = rb_const_get(PagloModule, rb_intern("PluginInstaller"));
  PluginInstaller = rb_funcall(PluginInstallerClass, rb_intern("new"), 1, PluginManager);

  /*
   * Tell the Ruby GC about the variables we will keep.
   */
  rb_gc_register_address(&PluginManager);
  rb_gc_register_address(&PluginInstaller);

  return PluginManager;
}
//---------------------------------------------------------------------------

bool TRubyThread::InitInterpreter(const char *Path)
{
  int Error;

  ruby_init();
  ruby_init_loadpath();
  
  PluginManager = rb_protect(ProtectedInitInterpreter, rb_str_new2(Path), &Error);
  if (Error) {
    HandleRubyError();
    return false;
  }

  return true;
}
//---------------------------------------------------------------------------

static VALUE ProtectedGetPlugins(VALUE Arg) 
{
  _PluginList.clear();

  VALUE rb_Plugins = rb_funcall(PluginManager, rb_intern("plugins"), 0);
  if (rb_Plugins != T_NIL) {
    
    int PluginCount = NUM2INT(rb_funcall(rb_Plugins, rb_intern("size"), 0));
    for (int i = 0; i < PluginCount; i++) {
      /*
       * Get the plugin class, name, and description.
       */
      VALUE rb_Plugin = rb_ary_pop(rb_Plugins);
      TPlugin Plugin;
      VALUE rb_Class = rb_hash_aref(rb_Plugin, rb_eval_string(":class"));
      Plugin.Class = StringValueCStr(rb_Class);
      VALUE rb_Name = rb_hash_aref(rb_Plugin, rb_eval_string(":name"));
      Plugin.Name = StringValueCStr(rb_Name);
      VALUE rb_Description = rb_hash_aref(rb_Plugin, rb_eval_string(":description"));
      Plugin.Description = StringValueCStr(rb_Description);
      
      /*
       * Unpack the plugin parameters.
       */
      VALUE rb_Params = rb_hash_aref(rb_Plugin, rb_eval_string(":params"));
      int ParamCount = NUM2INT(rb_funcall(rb_Params, rb_intern("size"), 0));
      for (int j = 0; j < ParamCount; j++) {
        /*
         * Get the parameter name.
         */
        VALUE rb_Param = rb_ary_pop(rb_Params);
        TPluginParam Param;
        rb_Name = rb_hash_aref(rb_Param, rb_eval_string(":name"));
        Param.Name = Plugin.Class + "." + StringValueCStr(rb_Name);
        
        /*
         * Unpack the options hash and load it into the 
         * options map.
         */
        VALUE rb_Options = rb_hash_aref(rb_Param, rb_eval_string(":options"));
        VALUE rb_Keys = rb_funcall(rb_Options, rb_intern("keys"), 0);
        int KeyCount = NUM2INT(rb_funcall(rb_Keys, rb_intern("size"), 0));
        for (int k = 0; k < KeyCount; k++) {
          VALUE rb_Key = rb_ary_pop(rb_Keys);
          VALUE rb_Value = rb_hash_aref(rb_Options, rb_Key);
          rb_Key = rb_funcall(rb_Key, rb_intern("id2name"), 0);
          string Key = StringValueCStr(rb_Key);
          string Value = StringValueCStr(rb_Value);
          Param.Options[Key] = Value;
        }
        Plugin.Params.push_back(Param);
      }
      _PluginList.push_back(Plugin);
    }
  }
  return rb_Plugins;
}
//---------------------------------------------------------------------------

bool TRubyThread::GetPlugins(void)
{
  int Error;
  rb_protect(ProtectedGetPlugins, 0, &Error);
  if (Error) {
    HandleRubyError();
    return false;
  }
  
  return true;
}
//---------------------------------------------------------------------------

static VALUE ProtectedExecuteDevicePlugins(VALUE NetDevHandle)
{
  VALUE Device = rb_funcall(NetworkDeviceClass, rb_intern("new"), 1, NetDevHandle);

  rb_funcall(PluginManager, rb_intern("execute_device_plugins"), 1, Device);

  return Qnil;
}
//---------------------------------------------------------------------------

TScanResult *TScanPlugin::PerformScan(TPThread *Thread)
{
  int Error;
  
  TNetworkDevice *NetDev = _AnalysisMgr->GetDeviceManager()->LookupDevice(this->TargetIP);
  rb_protect(ProtectedExecuteDevicePlugins, INT2NUM((long)NetDev), &Error);
  if (Error) {
    HandleRubyError();
  }
  rb_gc_start();
  
  return new TScanResult(this->TargetIP);
}
//---------------------------------------------------------------------------

static VALUE ProtectedInstallPlugin(VALUE PluginPath)
{
  rb_funcall(PluginInstaller, rb_intern("install_plugin"), 1, PluginPath);

  return Qnil;
}
//---------------------------------------------------------------------------

static VALUE ProtectedDownloadInstallPlugin(VALUE PluginURL)
{
  rb_funcall(PluginInstaller, rb_intern("download_and_install_plugin"), 1, PluginURL);

  return Qnil;
}
//---------------------------------------------------------------------------

TScanResult *TInstallPlugin::PerformScan(TPThread *Thread)
{
  int Error;
  string ErrorMessage = "";

  if (FromURL) {
    rb_protect(ProtectedDownloadInstallPlugin, rb_str_new2(Path.c_str()), &Error);
  } else {
    rb_protect(ProtectedInstallPlugin, rb_str_new2(Path.c_str()), &Error);
  }
  if (Error) {
    HandleRubyError();

    VALUE lasterr = rb_gv_get("$!");
    VALUE message = rb_obj_as_string(lasterr);

    ErrorMessage = RSTRING(message)->ptr;
  } else {
    /*
     * Reload the plugin list of the install was successful.
     */
    _RubyThread->GetPlugins();
  }
  
  return new TScanResult(ErrorMessage);
}
//---------------------------------------------------------------------------