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

#include "common.hh"

#include <fstream>
#include <memory>

#include "config.hh"
#include "buffer.hh"
#include "paths.hh"
#include "logger.hh"
#include "lru.hh"
#include "mitm.hh"
#include "ssl.hh"

namespace {

class ConnectionImpl : public Mitm::Connection {
public:
  ConnectionImpl(Logger* logger)
    : logger_(logger){
  }

  bool connect(SSLEntropy* entropy, SSLCertStore* store,
               std::string const& cert, SSLKey* key,
               bool unsecure, std::string const& host) {
    uint16_t flags = unsecure ? SSL::UNSECURE : 0;
    cert_.reset(SSLCert::load(logger_, cert));
    if (!cert_) return false;
    local_.reset(SSL::server(logger_, entropy, cert_.get(), key, flags));
    if (!local_) return false;
    remote_.reset(SSL::client(logger_, entropy, store, host, flags));
    if (!remote_) return false;

    in_.reset(Buffer::create(8192, 1024));
    out_.reset(Buffer::create(8192, 1024));

    return true;
  }

  bool transfer(Buffer* local_in, Buffer* local_out,
                Buffer* remote_in, Buffer* remote_out,
                Mitm::Monitor* monitor) override {
    bool local_active = true, remote_active = true;
    while (local_active || remote_active) {
      local_active = false;
      if (local_) {
        size_t avail;
        out_->read_ptr(&avail);
        switch (local_->transfer(local_in, local_out, in_.get(), out_.get())) {
        case SSL::NO_ERR: {
          size_t avail2;
          auto ptr = out_->read_ptr(&avail2);
          if (avail < avail2) {
            if (monitor) {
              monitor->local2remote(
                  static_cast<char const*>(ptr) + avail, avail2 - avail);
            }
            local_active = true;
          }
          break;
        }
        case SSL::ERR:
          return false;
        case SSL::CLOSED:
          local_.reset();
          break;
        }
      }
      remote_active = false;
      if (remote_) {
        size_t avail;
        in_->read_ptr(&avail);
        switch (remote_->transfer(remote_in, remote_out,
                                  out_.get(), in_.get())) {
        case SSL::NO_ERR: {
          size_t avail2;
          auto ptr = in_->read_ptr(&avail2);
          if (avail < avail2) {
            if (monitor) {
              monitor->remote2local(
                  static_cast<char const*>(ptr) + avail, avail2 - avail);
            }
            remote_active = true;
          }
          break;
        }
        case SSL::ERR:
          return false;
        case SSL::CLOSED:
          remote_.reset();
          break;
        }
      }
    }
    return true;
  }

  bool local_eof() const override {
    return !local_;
  }

  bool remote_eof() const override {
    return !remote_;
  }

  void close_local() override {
    if (local_) {
      local_->close();
    }
  }
  void close_remote() override {
    if (remote_) {
      remote_->close();
    }
  }

private:
  Logger* const logger_;
  std::unique_ptr<SSLCert> cert_;
  std::unique_ptr<SSL> local_;
  std::unique_ptr<SSL> remote_;
  std::unique_ptr<Buffer> in_;
  std::unique_ptr<Buffer> out_;
};

class MitmImpl : public Mitm {
public:
  MitmImpl(Logger* logger)
    : logger_(logger), unsecure_(false), cache_(42) {
  }

  ~MitmImpl() override {
  }

  bool load(Config* config, std::string const& cwd) {
    auto bundle_file = config->get("ssl_cert_bundle", nullptr);
    auto ca_cert_file = config->get("ssl_ca_cert", nullptr);
    auto ca_key_file = config->get("ssl_ca_key", nullptr);
    if (!bundle_file || !ca_cert_file || !ca_key_file) return false;
    entropy_.reset(SSLEntropy::create(logger_));
    if (!entropy_) return false;
    store_.reset(SSLCertStore::create(logger_, bundle_file));
    if (!store_) return false;
    std::string ca_cert, ca_key;
    if (!load_file(ca_cert_file, cwd, &ca_cert)) return false;
    if (!load_file(ca_key_file, cwd, &ca_key)) return false;
    unsecure_ = config->get("ssl_unsecure", false);
    issuer_cert_.reset(SSLCert::load(logger_, ca_cert));
    if (!issuer_cert_) return false;
    issuer_key_.reset(SSLKey::load(logger_, ca_key, entropy_.get()));
    if (!issuer_key_) return false;
    return true;
  }

  bool reload(Config* config, std::string const& cwd) override {
    std::unique_ptr<SSLCertStore> store(
        SSLCertStore::create(logger_, config->get("ssl_cert_bundle", "")));
    if (!store) return false;
    std::string ca_cert, ca_key;
    if (!load_file(config->get("ssl_ca_cert", ""), cwd, &ca_cert)) return false;
    if (!load_file(config->get("ssl_ca_key", ""), cwd, &ca_key)) return false;
    unsecure_ = config->get("ssl_unsecure", false);
    std::unique_ptr<SSLCert> issuer_cert(SSLCert::load(logger_, ca_cert));
    if (!issuer_cert) return false;
    std::unique_ptr<SSLKey> issuer_key(
        SSLKey::load(logger_, ca_key, entropy_.get()));
    if (!issuer_key) return false;
    store_.swap(store);
    issuer_cert_.swap(issuer_cert);
    issuer_key_.swap(issuer_key);
    return true;
  }

  DetectResult detect(void const* data, size_t avail) override {
    if (avail == 0) return NEED_MORE;
    auto d = static_cast<uint8_t const*>(data);
    // SSL 3.0
    if (d[0] == 0x16) {
      if (avail < 2) return NEED_MORE;
      if (d[1] != 0x03) return OTHER;  // Need fixing when SSL 4.0 shows up
      if (avail < 5) return NEED_MORE;
      if (((d[3] << 8) | d[4]) < 9)  // Min size of client hello
        return OTHER;
      if (avail < 6) return NEED_MORE;
      return d[5] == 0x01 ? SSL : OTHER;
    }
    // SSL 2.0
    if (!(d[0] & 0x80)) return OTHER;
    if (avail < 2) return NEED_MORE;
    if (((d[0] & 0x7f) << 8 | d[1]) < 9)  // Min size of client hello
      return OTHER;
    if (avail < 3) return NEED_MORE;
    return d[2] == 0x01 ? SSL : OTHER;
  }

  Connection* open(std::string const& host) override {
    CacheEntry* entry;
    CacheEntry _entry;
    entry = cache_.get(host);
    if (!entry) {
      entry = &_entry;
      if (!SSLCert::generate(logger_, entropy_.get(),
                             issuer_cert_.get(), issuer_key_.get(), host,
                             issuer_key_.get(), &entry->cert)) {
        return nullptr;
      }
      cache_.insert(host, _entry);
    }
    std::unique_ptr<ConnectionImpl> conn(new ConnectionImpl(logger_));
    if (!conn->connect(entropy_.get(), store_.get(), entry->cert,
                       issuer_key_.get(), unsecure_, host)) return nullptr;
    return conn.release();
  }

private:
  struct CacheEntry {
    std::string cert;
  };

  bool load_file(std::string const& path, std::string const& cwd,
                 std::string* out) const {
    std::ifstream in(Paths::join(cwd, path));
    if (!in.good()) {
      logger_->out(Logger::ERR, "Unable to open: %s", path.c_str());
      return false;
    }
    out->clear();
    while (in.good()) {
      char buffer[8192];
      in.read(buffer, sizeof(buffer));
      out->append(buffer, in.gcount());
    }
    if (!in.eof()) {
      logger_->out(Logger::ERR, "Unable to open: %s", path.c_str());
      return false;
    }
    return true;
  }

  Logger* const logger_;
  std::unique_ptr<SSLEntropy> entropy_;
  std::unique_ptr<SSLCertStore> store_;
  std::unique_ptr<SSLCert> issuer_cert_;
  std::unique_ptr<SSLKey> issuer_key_;
  bool unsecure_;
  lru<std::string, CacheEntry> cache_;
};

}  // namespace

// static
Mitm* Mitm::create(Logger* logger, Config* config, std::string const& cwd) {
  std::unique_ptr<MitmImpl> mitm(new MitmImpl(logger));
  if (!mitm->load(config, cwd)) return nullptr;
  return mitm.release();
}