Add basic console monitor and implement monitor support

1 files changed, 271 insertions, 0 deletions

diff --git a/src/monitor-cmd.cc b/src/monitor-cmd.cc
new file mode 100644
index 0000000..5221db0
--- /dev/null
+++ b/src/monitor-cmd.cc
@@ -0,0 +1,271 @@
+// -*- 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);
+  }
+}