// -*- mode: c++; c-basic-offset: 2; -*-

#include "common.hh"

#include <cstring>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <signal.h>
#include <unordered_map>

#include "args.hh"
#include "buffer.hh"
#include "io.hh"
#include "ios_save.hh"
#include "looper.hh"
#include "resolver.hh"
#include "monitor.hh"

namespace {

class Delegate : public Monitor::Delegate {
public:
  Delegate(std::ostream& out, bool interleave, Looper* looper)
    : out_(out), interleave_(interleave), looper_(looper), attached_(false) {
  }

  void state(Monitor* monitor, Monitor::State state) override {
    switch (state) {
    case Monitor::DISCONNECTED:
      std::cout << "# Disconnected" << std::endl;
      looper_->quit();
      break;
    case Monitor::CONNECTING:
      break;
    case Monitor::CONNECTED:
      if (attached_) {
        std::cout << "# Detached" << std::endl;
        attached_ = false;
      } else {
        std::cout << "# Connected" << std::endl;
        monitor->attach();
      }
      break;
    case Monitor::ATTACHED:
      std::cout << "# Attached" << std::endl;
      attached_ = true;
      break;
    }
  }

  void error(Monitor* UNUSED(monitor), std::string const& error) override {
    std::cerr << "# Error: " << error << std::endl;
  }

  void package(
      Monitor* UNUSED(monitor), Monitor::Package const& package) override {
    packages_.insert(std::make_pair(package.id, package));
  }

  void package_data(Monitor* UNUSED(monitor), uint32_t id,
                    char const* data, size_t size, bool last) override {
    auto it = packages_.find(id);
    if (it == packages_.end()) {
      assert(false);
      return;
    }
    if (interleave_) {
      print_package(it->second, last, data, size);
      if (last) packages_.erase(it);
      return;
    }
    auto buf = data_.find(id);
    if (last) {
      if (buf == data_.end()) {
        print_package(it->second, true, data, size);
      } else {
        buf->second->write(data, size);
        size_t avail;
        auto ptr = buf->second->read_ptr(&avail);
        print_package(it->second, true,
                      reinterpret_cast<char const*>(ptr), avail);
        data_.erase(buf);
      }
      packages_.erase(it);
    } else {
      if (buf == data_.end()) {
        buf = data_.insert(std::make_pair(id, Buffer::create(8192, 0))).first;
      }
      buf->second->write(data, size);
    }
  }

private:
  void print_package(
      Monitor::Package& pkg, bool last, char const* data, size_t size) {
    if (size == 0 && !last) return;
    {
      ios_save save(out_);
      out_ << "*** " << pkg.timestamp.tv_sec << '.'
           << std::setfill('0') << std::setw(9) << pkg.timestamp.tv_nsec
           << '\n';
    }
    out_ << "* Source: " << pkg.source_host << ':' << pkg.source_port << '\n'
         << "* Target: " << pkg.target_host << ':' << pkg.target_port << '\n';
    if (interleave_) {
      auto offset = offset_.find(pkg.id);
      if (offset != offset_.end()) {
        out_ << "* Bytes: " << offset->second << "-";
        if (last) {
          out_ << (offset->second + size);
          offset_.erase(offset);
        } else {
          offset->second += size;
        }
        out_ << '\n';
      } else if (!last) {
        offset_.insert(std::make_pair(pkg.id, size));
        out_ << "* Bytes: 0-";
      }
    }
    out_ << "* Size: " << size << '\n';
    if (size > 0) {
      ios_save save(out_);
      auto d = reinterpret_cast<uint8_t const*>(data);
      out_.flags(std::ios::hex);
      out_.fill('0');
      for (size_t i = 0; i < size; i += 16) {
        out_ << std::setw(8) << i << ' ';
        unsigned j = 0;
        for (; j < 8; ++j) {
          auto k = i + j;
          if (k >= size) break;
          out_ << ' ' << std::setw(2) << static_cast<int>(d[k]);
        }
        for (; j < 8; ++j) {
          out_ << "   ";
        }
        out_ << ' ';
        for (; j < 16; ++j) {
          auto k = i + j;
          if (k >= size) break;
          out_ << ' ' << std::setw(2) << static_cast<int>(d[k]);
        }
        for (; j < 16; ++j) {
          out_ << "   ";
        }
        out_ << "  |";
        j = 0;
        for (; j < 16; ++j) {
          auto k = i + j;
          if (k >= size) break;
          out_ << printable(data[k]);
        }
        for (; j < 16; ++j) {
          out_ << ' ';
        }
        out_ << "|\n";
      }
    }
    out_ << std::endl;
  }

  static char printable(char c) {
    return (c & 0x80 || c < ' ' || c >= 0x7f) ? '.' : c;
  }

  std::ostream& out_;
  bool interleave_;
  Looper* looper_;
  bool attached_;
  std::unordered_map<uint32_t, Monitor::Package> packages_;
  // Used when interleaving
  std::unordered_map<uint32_t, uint64_t> offset_;
  // Used when not interleaving
  std::unordered_map<uint32_t, std::unique_ptr<Buffer>> data_;
};

io::auto_pipe signal_pipe;

void signal(int UNUSED(signum)) {
  io::write_all(signal_pipe.write(), "", 1);
  std::cerr << "# Caught signal" << std::endl;
}

void quit_loop(Looper* looper, int UNUSED(fd), uint8_t UNUSED(events)) {
  looper->quit();
}

bool run(std::ostream& out, bool interleave,
         std::string const& host, uint16_t port) {
  std::unique_ptr<Looper> looper(Looper::create());
  std::unique_ptr<Resolver> resolver(Resolver::create(looper.get()));
  std::unique_ptr<Delegate> delegate(
      new Delegate(out, interleave, looper.get()));
  std::unique_ptr<Monitor> monitor(
      Monitor::create(looper.get(), resolver.get(), delegate.get()));
  std::cout << "# Connecting to " << host << ':' << port << std::endl;
  monitor->connect(host, port);
  if (signal_pipe.open()) {
    struct sigaction action;
    memset(&action, 0, sizeof(action));
    action.sa_handler = signal;
    sigaction(SIGINT, &action, nullptr);
    sigaction(SIGTERM, &action, nullptr);
    looper->add(signal_pipe.read(), Looper::EVENT_READ,
                std::bind(&quit_loop, looper.get(), std::placeholders::_1,
                          std::placeholders::_2));
  }
  return looper->run();
}

}  // namespace

int main(int argc, char** argv) {
  std::unique_ptr<Args> args(Args::create());
  args->add('i', "interleave",
            "unless set packages are not output until they are complete");
  args->add('o', "output", "FILE", "output packages to FILE instead of stdout");
  args->add('h', "help", "display this text and exit");
  args->add('V', "version", "display version and exit");
  if (!args->run(argc, argv)) {
    std::cerr << "Try `tp-monitor --help` for usage." << std::endl;
    return EXIT_FAILURE;
  }
  if (args->is_set('h')) {
    std::cout << "Usage: `tp-monitor [OPTIONS...] HOST[:PORT]`\n"
              << "Connect to transparent proxy running on HOST[:PORT] and "
              << "print logged packages\n"
              << std::endl;
    args->print_help();
    return EXIT_SUCCESS;
  }
  if (args->is_set('V')) {
    std::cout << "TransparentProxy monitor version " VERSION
              << " written by Joel Klinghed <the_jk@yahoo.com>" << std::endl;
    return EXIT_SUCCESS;
  }
  if (args->arguments().size() != 1) {
    std::cerr << "Unexpected number of arguments, expects one argument."
              << std::endl;
    return EXIT_FAILURE;
  }
  std::string host = args->arguments().front();
  uint16_t port;
  auto i = host.rfind(':');
  if (i != std::string::npos && i > 0 && host[i - 1] != ':') {
    errno = 0;
    char* end = nullptr;
    auto tmp = strtoul(host.c_str() + i + 1, &end, 10);
    if (errno || !end || *end) {
      std::cerr << "Invalid port number: " << host.substr(i + 1) << std::endl;
      return EXIT_FAILURE;
    }
    port = tmp & 0xffff;
  } else {
    port = 9000;
  }
  auto interleave = args->is_set('i');
  auto output = args->arg('o', nullptr);
  if (output) {
    std::ofstream out(output, std::ofstream::trunc);
    if (!out.good()) {
      std::cerr << "Unable to open " << output << " for writing." << std::endl;
      return EXIT_FAILURE;
    }
    return run(out, interleave, host, port);
  } else {
    return run(std::cout, interleave, host, port);
  }
}