//////////////////////////////////////////////////////////////////////
// IP traffic monitor
// Written by Jonathan A. Foster <ChipMaster@YeOlPiShack.net>
// Started August 11th, 2021
//
// The idea is to analyze iptables LOG entries in combination with
// DNSmasq's query log entries and combine them to list the hosts
// that were accessed. This will be read in real time from SysLogD
// using the "pipe" feature. Access records will be logged to MySQL
// for reporting and record keeping.
//////////////////////////////////////////////////////////////////////
// TODO: catch signals and attempt removal of PID file during shutdown.
#include <string.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <syslog.h>
#include <unistd.h>
#include <string>
#include <iostream>
#include <fstream>
#include <stdexcept>
#include <vector>
#include <map>

#include "../cli.h"
#include "../data.h"
#include "../config.h"
#include "appbase.h"
using namespace std;



//////////////////////////////////////////////////////////////////////
// Monitor Config
//////////////////////////////////////////////////////////////////////

struct MonitorConf: public MonitorBaseConf {
  INIusList    us;

  MonitorConf() { groups["us"] = &us; }
};



//////////////////////////////////////////////////////////////////////
// Application class to store data passed in through a pipe or
// file(s).
//////////////////////////////////////////////////////////////////////
//#define DEBUG

struct TrafficMon: public TrafficMonBaseApp {
  LogAnalyzer     analyze;
  istream        *log;
  LiveBug         bug;
  bool            setup_db;
  bool            piping;               // are we sucking on a pipe?
  bool            background_me;        // daemonize?
  string          pid_file;             // Name of file to write PID in
  int             inp_ct;               // How many inputs were on the CLI
  bool            running;              // Whether or not its time to run logs
  long long       line_no;              // current line # in the active input



  TrafficMon():
    log(              0),
    setup_db(     false),
    piping(       false),
    background_me(false),
    inp_ct(           0),
    running(      false),
    line_no(          0)
  {
    config = new MonitorConf;
    analyze.us = &(((MonitorConf *)config)->us.vals);
  }



  int help() {
    cerr <<
      "\n"
      "trafficmon -c {config file} [-i] [-d] [-b] [-p {pipe name} | [{log name} ...]]\n"
      "\n"
      "Arguments:\n"
      "  -c  Load configuration from file (required)\n"
      "  -i  PID file to write (optional, use with -b)\n"
      "  -d  Create data tables (optional, use once)\n"
      "  -b  Run in background - daemonize (optional, use -i too)\n"
      "  -p  read from pipe creating it if needed (optional, good with -b)\n"
      "\n"
      "Other arguments are log files to be processed. This can be used to bulk\n"
      "load data prior to going live with an always on daemon or for catching\n"
      "up if the daemon stopped for some reason.";
    return ExitCode = 1;
  }



  unsigned do_switch(const char *sw) {
    if(sw[1]==0) {
      switch(*sw) {
      case 'b': background_me = true; return 0;
      case 'd': setup_db      = true; return 0;
      case 'i':                       return 1;
      case 'p': piping        = true; return !running; // second pass treats it as a file
      }
    }
    return TrafficMonBaseApp::do_switch(sw);
  }



  void do_switch_arg(const char *sw, const string &val) {
    // If "running" is set then we've already done this. So skip it!
    if(running) return;

    if(sw[1]==0) switch(*sw) {

    /// PID file to write ID in if we're daemonizing ///

    case 'i':
      pid_file = val;
      return;

    /// Make pipe if requested ///

    case 'p':
      // This can't wait for "running" since we want it to fail in foreground
      // Make read/write by process owner (root) only
      if(mkfifo(val.c_str(), 0600)<0 && errno!=17 /*already exists*/) {
        ExitCode = 2;
        throw runtime_error("Making pipe raised error "+str(errno));
      }
      inp_ct++;
    }

    TrafficMonBaseApp::do_switch_arg(sw, val);
  }



  void do_arg(const char *fname) {
    // if not "running" we don't want to do this yet. Just count inputs.
    if(!running) {
      inp_ct++;
      return;
    }

    // the only thing we expect on the CLI is a log/pipe name
    if(log==&cin) log = 0; // won't happen right now
    if(log)
      ((ifstream*)log)->close();
    else
      log = new ifstream;
    if(!background_me) cerr << fname << ":\n";
restart:
    ((ifstream*)log)->open(fname);
    ExitCode = run();
    if(piping) {
      // If a process closes the write end of the pipe, like during log
      // rotation, we receive an EOF, end up here, and have to re-open the pipe
      // to be able to receive more on it.
      ((ifstream*)log)->close();
      sleep(1); // This is just to make sure we don't hog CPU in a loop
      goto restart;
    }
  }



  // This is the actual data process
  // probably should daemonize around here.
  int run() {
    string l;
    // prepare insert statement
    cppdb::statement db_add = db <<
      "INSERT INTO connections "
      "(us, us_port, them, them_port, them_name, protocol, inbound) "
      "VALUES (?,?,?,?,?,?,?)";

    /// parse log file ///

    line_no=0;
    while(getline(*log, l)) {
      line_no++;
      if(!background_me)
        cerr << bug << ' ' << line_no << '\r' << flush;

      /// process connections ///

      if(analyze.line(l)) {
        // TODO: time stamp handling?
        // insert record
        db_add.reset();
        db_add << analyze.conn.us << analyze.conn.us_port << analyze.conn.them
               << analyze.conn.them_port << analyze.conn.name
               << analyze.conn.protocol << analyze.conn.in
               << cppdb::exec;
      }
    }
    // In a real pipe situation this should never get reached.
    if(!background_me) {
      cerr << "\nLines:      " << line_no
           << "\nTotal rDNS: " << analyze.rdns.size() << '\n' << endl;
    }
    return 0;
  }



  void create_tables() {

    /// Connection recrods ///

    db << "CREATE TABLE connections ("
            "id        INT        NOT NULL AUTO_INCREMENT PRIMARY KEY,"
            "tstamp    TIMESTAMP  NOT NULL," // default is NOW()
            "us        CHAR(80)   NOT NULL DEFAULT '',"
            "us_port   INT        NOT NULL DEFAULT 0,"
            "them      CHAR(80)   NOT NULL DEFAULT '',"
            "them_port INT        NOT NULL DEFAULT 0,"
            "them_name CHAR(128)  NOT NULL DEFAULT '',"
            "protocol  CHAR(12)   NOT NULL DEFAULT '',"
            "inbound   TINYINT(1) NOT NULL DEFAULT 0,"
            "INDEX tstamp(tstamp, them_name)"
          ") Engine=MyISAM"
       << cppdb::exec;
    // NOTE: MyISAM is thousands of times faster than InnoDB for this sort
    //       of work load. Mostly continuous writes with occasional reads.
    //       But will it work for all I'm going to do? AriaDB is slower but
    //       significantly faster than InnoDB.

    /// Status of domain names ///

    db << "CREATE TABLE dns ("
            "name CHAR(128) NOT NULL PRIMARY KEY,"
            "decided TIMESTAMP NOT NULL,"
            "status TINYINT(1) NOT NULL DEFAULT 0,"
            "note CHAR(128) NOT NULL DEFAULT ''"
          ") Engine=MyISAM"
       << cppdb::exec;

    /// wild card DNS list ///

    // NOTE: All of these names are treated as if prefixed with "*.". At this
    //       time there are only plans to implement this as a black list.
    //       Many domains are sevices with threat potential and discovering
    //       all of the host and subdomains may not be viable. Think ad and
    //       tracking services.
    db << "CREATE TABLE dns_wild ("
            "name CHAR(128) NOT NULL PRIMARY KEY,"
            "decided TIMESTAMP NOT NULL,"
            "status TINYINT(1) NOT NULL DEFAULT 2,"
            "note CHAR(128) NOT NULL DEFAULT ''"
          ") Engine=MyISAM"
       << cppdb::exec;

  }



  int main() {
    int x;
    try {
      if(x=TrafficMonBaseApp::main()) return x;
      if(!((MonitorConf*)config)->us.vals.size()) throw CLIerror(
        "The configuration files MUST contain an [us] section with "
        "appropriate values"
      );
      if(!config->traffic_mon.vals.size()) throw CLIerror(
        "The configuration files MUST contain an [Traffic Mon] section with "
        "appropriate values"
      );
      if(piping && inp_ct!=1) throw CLIerror(
        "Pipe requires one and only one file name to read from."
      );
      if(setup_db) create_tables();

      // if requested attempt daemonizing.
      if(background_me) {
        if(inp_ct==0) throw CLIerror(
          "Backgrounding requires the specification of an input source. "
          "STDIN is unavailable when we move to background."
        );

        if(daemon(0,0)<0) throw CLIerror("Failed to switch to background");
        // If we get here we're in the background. No STDIN, STDOUT, STDERR
        if(pid_file!="") { // pid_file requested
          ofstream pf(pid_file.c_str());
          pf << getpid();
        }
      }
    } catch(const CLIerror &e) {
      cerr << "ERROR: " << e.what() << "\n";
      help();
      return ExitCode ? ExitCode : 1; // JIC someone sets a different exit code
    }

    try {
      // CLI processed now lets analyze data.
      running = true;
// TODO: this is going to break things? Probably should prevent reloading conf. DB is opened in TrafficMonBaseApp::main()
// Actually I think it may actualy work...
      cBaseApp::main(); // re-run CLI args for inputs
      if(!log) {
        // no inputs spec'd on CLI assume stdin (we're not a daemon)
        log = &cin;
        ExitCode = run();
      }
    } catch(const exception &e) {
      if(ExitCode==0) ExitCode = 2;
      if(background_me) {
        openlog("trafficmon", LOG_CONS | LOG_PID, LOG_DAEMON);
        syslog(LOG_ALERT, "%s", e.what());
        closelog();
      } else
        cerr << "ERROR: " << e.what() << "\n";
    }
    return ExitCode;
  }

};



//////////////////////////////////////////////////////////////////////
// Run it
//////////////////////////////////////////////////////////////////////

MAIN(TrafficMon)