twproxy2/session.cpp

#include <boost/bind.hpp>
#include <iostream>

#include <boost/log/attributes/named_scope.hpp>
#include <boost/log/core.hpp>
#include <boost/log/trivial.hpp>

#include <regex>

#include "session.h"

#include <string>

bool replace(std::string &str, const std::string &from, const std::string &to) {
  size_t start_pos = str.find(from);
  if (start_pos == std::string::npos)
    return false;
  do {
    str.replace(start_pos, from.length(), to);
  } while ((start_pos = str.find(from)) != std::string::npos);
  return true;
}

bool replace(std::string &str, const char *from, const char *to) {
  size_t start_pos = str.find(from);
  if (start_pos == std::string::npos)
    return false;
  do {
    str.replace(start_pos, strlen(from), to);
  } while ((start_pos = str.find(from)) != std::string::npos);
  return true;
}

void ansi_clean(std::string &str) {
  static std::regex ansi_cleaner("\x1b\[[0-9;]*[A-Zmh]", std::regex_constants::ECMAScript);
  str = std::regex_replace(str, ansi_cleaner, "");
}

std::string clean_string(const std::string &source) {
  // BOOST_LOG_NAMED_SCOPE("clean_string");
  std::string clean = source;
  replace(clean, "\n", "\\n");
  replace(clean, "\r", "\\r");

  // ANSI too
  ansi_clean(clean);
  // BOOST_LOG_TRIVIAL(error) << "cleaned: " << clean;

  replace(clean, "\x1b", "^");

  return clean;
}

session::session(boost::asio::ip::tcp::socket socket,
                 boost::asio::io_service &io_service, std::string hostname,
                 std::string port)
    : socket_(std::move(socket)),
      io_service_{io_service}, resolver_{io_service}, server_{io_service},
      timer_{io_service}, host{hostname}, port{port} {
  server_sent = 0;
  // BOOST_LOG_NAMED_SCOPE("session");
}

void session::start(void) {
  BOOST_LOG_NAMED_SCOPE("session");
  BOOST_LOG_TRIVIAL(info) << "session";
  auto self(shared_from_this());
  // read_buffer.reserve(1024);
  // do_write("Welcome!\n");
  do_read();
}

session::~session() {
  BOOST_LOG_NAMED_SCOPE("session");
  BOOST_LOG_TRIVIAL(info) << "~session destructed";
}

void session::parse_auth(void) {
  // how many nulls should I be seeing?
  // \0user\0pass\0terminal/SPEED\0
  // Maybe in the future I'll care about parsing this out.  I don't right now.

  // Ok, yes I do!  If I don't have a proper rlogin value here, it isn't going
  // to work when I try to connect to the rlogin server.

  if (rlogin_auth.size() > 10)
    rlogin_name = rlogin_auth.c_str() + 1;
  else
    rlogin_name = "?";
}

void session::on_connect(const boost::system::error_code error) {
  // We've connected to the server!  WOOT WOOT!
  BOOST_LOG_NAMED_SCOPE("session");

  if (!error) {
    BOOST_LOG_TRIVIAL(info) << "Connected to " << host;
    to_client("Connected...\n\r");
    connected = true;
    if (rlogin_auth[0] != 0) {
      // Ok, the rlogin information was junk --
      to_client("Let me make up some fake rlogin data for you...\n\r");
      char temp[] = "\0test\0test\0terminal/9600\0";
      std::string tmp(temp, sizeof(temp));
      to_server(tmp);
    } else {
      to_server(rlogin_auth);
    }

    server_read();
  } else {
    // TODO:
    std::string output = "Failed to connect : ";
    output += host;
    output += " : ";
    output += port;
    output += "\n\r";
    to_client(output);
    BOOST_LOG_TRIVIAL(error) << "Failed to connec to " << host << ":" << port;

    std::cout << "SHUTDOWN..." << std::endl;
    socket_.shutdown(boost::asio::ip::tcp::socket::shutdown_both);
  }
}

void session::dispatch_line(std::string line) {
  // Does this have \n\r still on it?  I don't want them.
  BOOST_LOG_NAMED_SCOPE("session");
  std::string temp = clean_string(line);
  BOOST_LOG_TRIVIAL(info) << "SL: " << temp; // clean_string(line);
  // std::cout << "SL: " << line << std::endl;
  // is echo on?  if so:
}

void session::process_lines(void) {
  // break server_prompt into lines and send/process one by one.
  size_t pos;
  // while ((pos = server_prompt.find("\n\r", 0, 2)) != std::string::npos) {
  // while ((pos = server_prompt.find("\r\n", 0, 2)) != std::string::npos) {
  while ((pos = server_prompt.find('\n', 0)) != std::string::npos) {
    // line
    std::string line = server_prompt.substr(0, pos + 1);
    server_prompt = server_prompt.substr(pos + 1);

    // Remove \n for dispatching
    std::string part = line.substr(0, pos);
    if (server_sent != 0) {
      line = line.substr(server_sent);
      server_sent = 0;
    };
    // display on?
    to_client(line);
    replace(part, "\r", "");
    dispatch_line(part);
  }

  // display on?
  if (server_sent != 0) {
    // send partial
    std::string part = server_prompt.substr(server_sent);
    to_client(part);
    server_sent = server_prompt.size();
  } else {
    // send all
    if (!server_prompt.empty()) {
      to_client(server_prompt);
      server_sent = server_prompt.size();
    }
  }

  // server_sent is the # of chars we've already sent of this.
}
void session::server_read(void) {
  auto self(shared_from_this());

  boost::asio::async_read(
      server_, boost::asio::buffer(server_buffer, sizeof(server_buffer) - 1),
      boost::asio::transfer_at_least(1),
      [this, self](boost::system::error_code ec, std::size_t length) {
        if (!ec) {
          server_buffer[length] = 0;

          server_prompt.append(server_buffer, length);
          process_lines();

          /*
            if (length) {
              // std::cout << length << std::endl;
              std::cout << "S: " << server_buffer << std::endl;
              do_write(server_buffer);
            }
            */

          server_read();
        } else {
          std::cout << "S: read_failed: connection closed" << std::endl;
          socket_.shutdown(boost::asio::ip::tcp::socket::shutdown_both);
          // socket_.async_shutdown(boost::bind(&session::on_shutdown, this,
          // boost::asio::placeholders::error));
        }
      });
}

void session::on_resolve(
    const boost::system::error_code error,
    const boost::asio::ip::tcp::resolver::results_type results) {
  //
  auto self(shared_from_this());

  if (!error) {
    // Take the first endpoint.
    boost::asio::ip::tcp::endpoint const &endpoint = *results;

    server_.async_connect(endpoint,
                          boost::bind(&session::on_connect, this,
                                      boost::asio::placeholders::error));

  } else {
    // TO DO:
    BOOST_LOG_NAMED_SCOPE("session");
    BOOST_LOG_TRIVIAL(error) << "Unable to resolve: " << host;
    std::string output = "Unable to resolve: ";
    output += host;
    output += "\n\r";
    to_client(output);
    std::cout << "SHUTDOWN ..." << std::endl;
    socket_.shutdown(boost::asio::ip::tcp::socket::shutdown_both);
  }
}

void session::do_read(void) {
  auto self(shared_from_this());

  boost::asio::async_read( // why can't I async_read_some here?
      socket_, boost::asio::buffer(read_buffer, sizeof(read_buffer) - 1),
      boost::asio::transfer_at_least(1),
      [this, self](boost::system::error_code ec, std::size_t length) {
        if (!ec) {
          read_buffer[length] = 0;
          if (rlogin_auth.empty()) {
            // first read should be rlogin information
            rlogin_auth.assign(read_buffer, length);

            // parse authentication information
            parse_auth();
            to_client(std::string(1, 0));

            to_client("Welcome, ");
            to_client(rlogin_name);
            to_client("\n\r");

            // Activate the connection to the server

            /* // this fails, and I'm not sure why.  I've used code like this
               before. resolver_.async_resolve( host, port, std::bind(
               &session::on_resolve, this, _1, _2)); */

            // This example shows using boost::bind, which WORKS.
            // https://stackoverflow.com/questions/6025471/bind-resolve-handler-to-resolver-async-resolve-using-boostasio
            resolver_.async_resolve(
                host, port,
                boost::bind(&session::on_resolve, this,
                            boost::asio::placeholders::error,
                            boost::asio::placeholders::iterator));

          } else if (length) {
            // std::cout << length << std::endl;

            // Proxy Active?
            BOOST_LOG_NAMED_SCOPE("session");

            to_server(read_buffer);
            BOOST_LOG_TRIVIAL(info) << "C: " << read_buffer;

            // do_write(output);
          }
          do_read();
        } else {
          std::cout << "C: read_failed: connection closed" << std::endl;
          if (connected)
            server_.shutdown(boost::asio::ip::tcp::socket::shutdown_both);
          // server_.async_shutdown(boost::bind(&session::on_shutdown, this,
          // boost::asio::placeholders::error));
        }
      });
}

void session::to_client(std::string message) {
  auto self(shared_from_this());
  boost::asio::async_write(
      socket_, boost::asio::buffer(message),
      [this, self](boost::system::error_code ec, std::size_t /*length*/) {
        if (!ec) {

        } else {
          std::cout << "write failed? closed?" << std::endl;
          server_.shutdown(boost::asio::ip::tcp::socket::shutdown_both);
          // server_.async_shutdown(boost::bind(&session::on_shutdown, this,
          // boost::asio::placeholders::error));
        }
      });
}

void session::to_server(std::string message) {
  auto self(shared_from_this());
  boost::asio::async_write(
      server_, boost::asio::buffer(message),
      [this, self](boost::system::error_code ec, std::size_t /*length*/) {
        if (!ec) {

        } else {
          std::cout << "write failed? closed?" << std::endl;
          socket_.shutdown(boost::asio::ip::tcp::socket::shutdown_both);
          // socket_.async_shutdown(boost::bind(&session::on_shutdown, this,
          // boost::asio::placeholders::error));
        }
      });
}

void session::on_shutdown(boost::system::error_code ec) {
  std::cout << "shutdown." << std::endl;
}

server::server(boost::asio::io_service &io_service,
               const boost::asio::ip::tcp::endpoint &endpoint, std::string host,
               std::string port)
    : io_service_{io_service}, acceptor_{io_service_, endpoint}, host_{host},
      port_{port} {
  do_accept();
}

/**
 * setup async connect accept
 *
 * This creates a session for each connection.  Using make_shared allows the
 * session to automatically clean up when it is no longer active/has anything
 * running in the reactor.
 */
void server::do_accept(void) {
  acceptor_.async_accept([this](boost::system::error_code ec,
                                boost::asio::ip::tcp::socket socket) {
    if (!ec) {
      std::make_shared<session>(std::move(socket), io_service_, host_, port_)
          ->start();
    }

    do_accept();
  });
}