/*
 *  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
 */
//---------------------------------------------------------------------------
#ifdef HAVE_CONFIG_H
#  include "config.h"
#endif
#include "ScanAgentAPI.h"
#include "AnalysisManager.h"
#include "NetworkDevice.h"
#include "EvidenceDescription.h"
#include "AgentConfig.h"
#include "ScanThread.h"
#include "ScanAgent.h"
#include "Debug.h"
#include "ServerInterface.h"
#include "CommunicationThread.h"
#include "SubmissionCache.h"

#ifndef PERSISTENT_DB
#  ifdef CONFIG_PERSISTENT_DB
#    define PERSISTENT_DB CONFIG_PERSISTENT_DB
#  else
#    define PERSISTENT_DB "persistent_data.db"
#  endif
#endif
#ifndef CONFIG_SUBMISSION_CACHE_DB
#  define CONFIG_SUBMISSION_CACHE_DB "submission_cache.db"
#endif

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

extern pthread_mutex_t _ScanThreadPoolMutex;
extern time_t _StartTime;


/*
 * Statistics about the scanning process.
 */
TScanningStatistics _Stats;

TDataStore _DataStore(PERSISTENT_DB);

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

int DeviceCount(void)
{
  TDeviceManager *DevMgr = _AnalysisMgr->GetDeviceManager();

  return DevMgr->Count();
}
//---------------------------------------------------------------------------

TDeviceHandle *GetDevice(int Index)
{
  TDeviceManager *DevMgr = _AnalysisMgr->GetDeviceManager();

  return DevMgr->LookupDevice(Index);
}
//---------------------------------------------------------------------------

void GetDeviceDetail(TDeviceHandle *Device, TDeviceDetail *Detail)
{
  TNetworkDevice *NetDev = (TNetworkDevice*)Device;

  Detail->IPAddress = NetDev->GetIPAddress().Print();
  Detail->MacAddress = NetDev->GetMacAddress().Print();

  if (NetDev->ClassificationCount() > 0) {
    TClassification C = NetDev->GetClassification(0);
    Detail->Vendor = C.GetValue(CR_VENDOR);
    Detail->Model = C.GetValue(CR_MODEL);
    Detail->Classes = C.GetValue(CR_CLASS);
    Detail->Score = C.Score;
  } else {
    Detail->Score = 0;
  }
}
//---------------------------------------------------------------------------

int IssueCount(void)
{
  return _IssuesList.size();
}
//---------------------------------------------------------------------------

TIssueInfo *GetIssue(int Index)
{
  if (Index < (int)_IssuesList.size()) {
    return _IssuesList[Index];
  } else {
    return NULL;
  }
}
//---------------------------------------------------------------------------

void ClearIssue(int Index)
{
  int Pos = 0;
  for (std::vector<TIssueInfo*>::iterator Iter = _IssuesList.begin();
    Iter != _IssuesList.end(); Iter++, Pos++) {
    if (Pos == Index) {
      _IssuesList.erase(Iter);
      return;
    }
  }
}
//---------------------------------------------------------------------------

void GetScanningStatistics(TScanningStatistics *Stats,
  bool IncludeEvidenceTypeCounts)
{
  /*
   * If we aren't running then reset the stats.
   */
  if (!_AnalysisMgr) {
    Stats->Running = false;
    return;
  }

  Stats->Running = true;
  Stats->AddressesScanned = _Stats.AddressesScanned;
  Stats->DevicesFound = DeviceCount();

  _ScanReqQueue.Lock();
  Stats->PendingProbes = _ScanReqQueue.Size();

  for (int i = 0; i < _ScanReqQueue.Size(); i++) {
    TPEvent *E = _ScanReqQueue.Get(i);
    if (E->GetType() == EVT_SCAN_REQUEST) {
      TScanRequest *Req = E->GetScanReq();
      if (Req) {
        Stats->PendingScanList.push_back(TProbeInfo(Req->GetID(), Req->GetType(),
         Req->Description(), Req->StartTime));
      }
    } else {
      // int a = 1;
      // TODO: Is this possible?
    }
  }

  _ScanReqQueue.Unlock();

  /*
   * List the threads.
   */
  Stats->ThreadsActive = 0;
  Stats->ActiveThreadList.clear();
  pthread_mutex_lock(&_ScanThreadPoolMutex);
  for(unsigned int i = 0; i < _ScanThreadPool.size(); i++) {
    TScanThread *Thread = _ScanThreadPool[i];
    if (Thread->GetThreadStatus() != "Idle") {
      TScanRequest *Req = Thread->GetScanReq();
      if (Req) {
        Stats->ThreadsActive++;
        Stats->ActiveThreadList.push_back(TProbeInfo(Req->GetID(),
          Req->GetType(), Req->Description(), Req->StartTime));
      }
    }
  }
  pthread_mutex_unlock(&_ScanThreadPoolMutex);

  /*
   * Add an entry for the ruby thread if it is active.
   */
  if (_RubyThread->Active) {
    Stats->ActiveThreadList.push_back(TProbeInfo(999999999,
      SCAN_PLUGINS, "Running plugin: " + _RubyThread->CurrentPluginName, _RubyThread->PluginStartTime));
  }

  Stats->CompletedProbes = _Stats.CompletedProbes;

  Stats->Submissions = _Submissions;
  Stats->Runtime = time(NULL) - _StartTime;
  Stats->CommBackoffSeconds = _CommThread->GetBackoffSeconds();

  Stats->FramesReceived = _SniffThread->GetReceived();
  Stats->FramesDropped = _SniffThread->GetDropped();

  /*
   * Queues.
   */
  Stats->CommunicationInputQueueSize = _CommThread->GetInputQueue()->Size();
  Stats->CommunicationFailedCmdSize = _CommThread->FailedCommandsSize();
  Stats->RubyInputQueueSize = _RubyThread->GetInputQueue()->Size();
  Stats->ScanReqQueueSize = _ScanReqQueue.Size();
  Stats->AnalysisInputQueueSize = _AnalysisMgr->GetInputQueue()->Size();
  Stats->AnalysisOutputQueueSize = _AnalysisMgr->GetOutputQueue()->Size();
  Stats->AnalysisPktQueueSize = _AnalysisMgr->GetPacketQueue()->Size();
  Stats->NumPacketQueues = _SniffThread->CountPacketQueues();

  /*
   * Count of evidence by type.
   */
  Stats->NumEvidenceTypes = MAX_MULTIMAP_TYPE;
  for (int i = 0; i < MAX_MULTIMAP_TYPE; i++) {
    Stats->EvidenceTypeCounts[i] = 0;
  }
  if (IncludeEvidenceTypeCounts) {
    Stats->EvidenceTypeNames[TYPE_NONE] = "None";
    Stats->EvidenceTypeNames[TYPE_IPADDRESS] = "IP address";
    Stats->EvidenceTypeNames[TYPE_MACADDRESS] = "MAC address";
    Stats->EvidenceTypeNames[TYPE_STRING] = "String";
    Stats->EvidenceTypeNames[TYPE_INTEGER] = "Integer";
    Stats->EvidenceTypeNames[TYPE_BOOL] = "Boolean";
    Stats->EvidenceTypeNames[TYPE_PROBEDATA] = "Probe data";
    Stats->EvidenceTypeNames[TYPE_SWITCHDATA] = "Switch data";
    Stats->EvidenceTypeNames[TYPE_ATTRIBUTEDATA] = "Attribute data";
    Stats->EvidenceTypeNames[TYPE_DHCPDATA] = "DHCP data";
    Stats->EvidenceTypeNames[TYPE_DISCOVERY_PROTO_DATA] = "Discovery data";
    Stats->EvidenceTypeNames[TYPE_WMI_DATA] = "WMI data";

    TDeviceManager *DevMgr = _AnalysisMgr->GetDeviceManager();
    for (int i = 0; i < DevMgr->Count(); i++) {
      TNetworkDevice *NetDev = DevMgr->LookupDevice(i);
      if (!NetDev) {
        break;
      }

      TMultiTypeMap Evidence = NetDev->StartDumpEvidence();
      try {
        for (TMultiTypeMap::iterator MapIterator = Evidence.begin();
          MapIterator != Evidence.end(); MapIterator++) {

          string Characteristic = (*MapIterator).first;
          TMultiTypeMapValue Value = (*MapIterator).second;

          Stats->EvidenceTypeCounts[Value.ValueType]++;
        }
        NetDev->EndDumpEvidence();
      } catch (exception &E) {
         NetDev->EndDumpEvidence();
        throw E;
      }
    }
  }

  /*
   * See when next network scan starts.
   */
  Stats->NextNetworkScanAt = _SchedulerThread->NextTaskTime(TASK_SCANNETWORK);
}
//---------------------------------------------------------------------------

void ClassificationCorrect(TDeviceHandle *Device)
{
  TNetworkDevice *NetDev = (TNetworkDevice*)Device;

  /*
   * Add a record indicating correct classification.
   */
  NetDev->AddEvidence(CT_CLASSIFICATION_FEEDBACK, string(RV_CORRECT));

  /*
   * Perform a transaction with the server to let it know the new classification.
   */
  _AnalysisMgr->AnalyzeDevice(NetDev);
}
//---------------------------------------------------------------------------

void ClassificationUpdate(TDeviceHandle *Device, const char *Vendor, const char *Model)
{
  TNetworkDevice *NetDev = (TNetworkDevice*)Device;

  /*
   * Add evidence records to the device for the new classification.
   */
  NetDev->AddEvidence(CT_ACTUAL_VENDOR, string(Vendor));
  NetDev->AddEvidence(CT_ACTUAL_MODEL, string(Model));

  /*
   * Perform a transaction with the server to let it know the new classification.
   */
  _AnalysisMgr->AnalyzeDevice(NetDev);
}
//---------------------------------------------------------------------------

void StartCrawlNow()
{
  _SchedulerThread->RunNow(TASK_SCANNETWORK);
}
//---------------------------------------------------------------------------

void DeleteDevices()
{
  _AnalysisMgr->GetDeviceManager()->Flush();
}
//---------------------------------------------------------------------------

void FlushSubmissionCache()
{
  TSubmissionCache SubmissionCache(CONFIG_SUBMISSION_CACHE_DB);
  if (SubmissionCache.Flush()) {
    /*
     * Don't delete the in-memory state because there are too many thread
     * related issues caused by doing it.
     * _AnalysisMgr->GetDeviceManager()->Flush();
     */
  }
  /*
   * Start a new scan to rediscover everything.
   */
  _SchedulerThread->RunNow(TASK_SCANNETWORK);
}
//---------------------------------------------------------------------------

void DumpMemory()
{
}
//---------------------------------------------------------------------------

std::string InstallPlugin(std::string Path, bool FromURL)
{
  TInstallPlugin *Scan = new TInstallPlugin(Path, FromURL);

  TEventQueue ReplyQueue;

  Scan->SetReplyQueue(&ReplyQueue);

  TPEvent *OutEvent = new TPEvent(Scan);
  _RubyThread->GetInputQueue()->Put(OutEvent);

  while (ReplyQueue.IsEmpty()) {
    MillisecondSleep(100);
  }

  TPEvent *ReplyEvent = ReplyQueue.Get();
  string Response = ReplyEvent->GetScanRes()->GetResponse();

  return Response;
}
//---------------------------------------------------------------------------

std::string ImmediateScanDevice(std::string IPAddrStr)
{
  try {
    TIPAddress IPAddr(IPAddrStr);
  
    TDeviceManager *DevMgr = _AnalysisMgr->GetDeviceManager();
    TNetworkDevice *NetDev = DevMgr->LookupDevice(IPAddr);

    if (NetDev) {
      /*
       * See if a scan is in progress.
       */
      if (NetDev->NumPendingScans() > 0) {
        return "A scan is already in progress for the device";
      } else {
        _AnalysisMgr->IssueScans(NetDev, false);
        return "Scan initiated";
      }
    } else {
      /*
       * The device is new, create a record for it.
       */
      NetDev = _AnalysisMgr->LookupOrCreateDevice(IPAddr);
      if (NetDev) {
        return "Device created and scan initiated";
      } else if (!_GlobalConfig.IsHomeNet(IPAddr)) {
        return "IP address lies outside of ranges that can be scanned";
      }
    }
  } catch (exception &E) {
    return E.what();
  }

  return "Failed to start scan";
}
//---------------------------------------------------------------------------