firegex-traffic-viewer/backend/proxy/proxy.cpp

#include <cstdlib>
#include <cstddef>
#include <iostream>
#include <string>
#include <csignal>
#include <fstream>
#include <regex>
#include <mutex>

#include <boost/thread.hpp>
#include <boost/shared_ptr.hpp>
#include <boost/enable_shared_from_this.hpp>
#include <boost/bind/bind.hpp>
#include <boost/asio.hpp>
#include <boost/thread/mutex.hpp>

//#define MULTI_THREAD
//#define DEBUG
//#define DEBUG_PACKET
//#define THREAD_NUM
using namespace std;

boost::asio::io_service *ios_loop = nullptr;

bool unhexlify(string const &hex, string &newString) {
   try{
      int len = hex.length();
      for(int i=0; i< len; i+=2)
      {
         std::string byte = hex.substr(i,2);
         char chr = (char) (int)strtol(byte.c_str(), NULL, 16);
         newString.push_back(chr);
      }
      return true;
   }
   catch (...){
      return false;
   }
}
struct regex_rules{
   vector<pair<string,regex>> regex_s_c_w, regex_c_s_w, regex_s_c_b, regex_c_s_b;
};
shared_ptr<regex_rules> regex_config;

const char* config_file;
mutex update_mutex;
#ifdef MULTI_THREAD
mutex stdout_mutex;
#endif 

bool filter_data(unsigned char* data, const size_t& bytes_transferred, vector<pair<string,regex>> const &blacklist, vector<pair<string,regex>> const &whitelist){
   #ifdef DEBUG_PACKET
   cerr << "---------------- Packet ----------------" << endl;
   for(int i=0;i<bytes_transferred;i++){
      cerr << data[i];
   }
   cerr << "\n" << "---------------- End Packet ----------------" << endl;
   #endif
   for (pair<string,regex> ele:blacklist){
      try{
         if(regex_search(reinterpret_cast<const char*>(data), reinterpret_cast<const char*>(data)+bytes_transferred, ele.second)){
            #ifdef MULTI_THREAD
            std::unique_lock<std::mutex> lck(stdout_mutex);
            #endif
            cout << "BLOCKED " << ele.first << endl;
            return false;
         }
      } catch(...){
         cerr << "Error while matching regex: " << ele.first << endl;
      }
   }
   for (pair<string,regex> ele:whitelist){
      try{
         if(!regex_search(reinterpret_cast<const char*>(data),reinterpret_cast<const char*>(data)+bytes_transferred, ele.second)){
            #ifdef MULTI_THREAD
            std::unique_lock<std::mutex> lck(stdout_mutex);
            #endif
            cout << "BLOCKED " << ele.first << endl;
            return false;
         }
      } catch(...){
         cerr << "Error while matching regex: " << ele.first << endl;
      }      
   }
   #ifdef DEBUG
   cerr << "Packet Accepted!" << endl;
   #endif
   return true;
}

namespace tcp_proxy
{
   namespace ip = boost::asio::ip;

   class bridge : public boost::enable_shared_from_this<bridge>
   {
   public:

      typedef ip::tcp::socket socket_type;
      typedef boost::shared_ptr<bridge> ptr_type;

      bridge(boost::asio::io_service& ios)
      : downstream_socket_(ios),
        upstream_socket_  (ios)
      {}

      socket_type& downstream_socket()
      {
         // Client socket
         return downstream_socket_;
      }

      socket_type& upstream_socket()
      {
         // Remote server socket
         return upstream_socket_;
      }

      void start(const string& upstream_host, unsigned short upstream_port)
      {
         // Attempt connection to remote server (upstream side)
         upstream_socket_.async_connect(
              ip::tcp::endpoint(
                   boost::asio::ip::address::from_string(upstream_host),
                   upstream_port),
               boost::bind(&bridge::handle_upstream_connect,
                    shared_from_this(),
                    boost::asio::placeholders::error));
      }

      void handle_upstream_connect(const boost::system::error_code& error)
      {
         if (!error)
         {
            // Setup async read from remote server (upstream)
            upstream_socket_.async_read_some(
                 boost::asio::buffer(upstream_data_,max_data_length),
                 boost::bind(&bridge::handle_upstream_read,
                      shared_from_this(),
                      boost::asio::placeholders::error,
                      boost::asio::placeholders::bytes_transferred));

            // Setup async read from client (downstream)
            downstream_socket_.async_read_some(
                 boost::asio::buffer(downstream_data_,max_data_length),
                 boost::bind(&bridge::handle_downstream_read,
                      shared_from_this(),
                      boost::asio::placeholders::error,
                      boost::asio::placeholders::bytes_transferred));
         }
         else
            close();
      }

   private:

      /*
         Section A: Remote Server --> Proxy --> Client
         Process data recieved from remote sever then send to client.
      */

      // Read from remote server complete, now send data to client
      void handle_upstream_read(const boost::system::error_code& error,
                                const size_t& bytes_transferred) // Da Server a Client
      {
         if (!error)
         {
            shared_ptr<regex_rules> regex_old_config = regex_config;
            if (filter_data(upstream_data_, bytes_transferred, regex_old_config->regex_s_c_b, regex_old_config->regex_s_c_w)){
               async_write(downstream_socket_,
                  boost::asio::buffer(upstream_data_,bytes_transferred),
                  boost::bind(&bridge::handle_downstream_write,
                        shared_from_this(),
                        boost::asio::placeholders::error));
            }else{
               close();
            }
         }
         else
            close();
      }

      // Write to client complete, Async read from remote server
      void handle_downstream_write(const boost::system::error_code& error)
      {
         if (!error)
         {

            upstream_socket_.async_read_some(
                 boost::asio::buffer(upstream_data_,max_data_length),
                 boost::bind(&bridge::handle_upstream_read,
                      shared_from_this(),
                      boost::asio::placeholders::error,
                      boost::asio::placeholders::bytes_transferred));
         }
         else
            close();
      }
      // *** End Of Section A ***


      /*
         Section B: Client --> Proxy --> Remove Server
         Process data recieved from client then write to remove server.
      */

      // Read from client complete, now send data to remote server
      void handle_downstream_read(const boost::system::error_code& error,
                                  const size_t& bytes_transferred) // Da Client a Server
      {
         if (!error)
         {
            shared_ptr<regex_rules> regex_old_config = regex_config;
            if (filter_data(downstream_data_, bytes_transferred, regex_old_config->regex_c_s_b, regex_old_config->regex_c_s_w)){
               async_write(upstream_socket_,
                  boost::asio::buffer(downstream_data_,bytes_transferred),
                  boost::bind(&bridge::handle_upstream_write,
                        shared_from_this(),
                        boost::asio::placeholders::error));
            }else{
               close();
            }
         }
         else
            close();
      }

      // Write to remote server complete, Async read from client
      void handle_upstream_write(const boost::system::error_code& error)
      {
         if (!error)
         {
            downstream_socket_.async_read_some(
                 boost::asio::buffer(downstream_data_,max_data_length),
                 boost::bind(&bridge::handle_downstream_read,
                      shared_from_this(),
                      boost::asio::placeholders::error,
                      boost::asio::placeholders::bytes_transferred));
         }
         else
            close();
      }
      // *** End Of Section B ***

      void close()
      {
         boost::mutex::scoped_lock lock(mutex_);

         if (downstream_socket_.is_open())
         {
            downstream_socket_.close();
         }

         if (upstream_socket_.is_open())
         {
            upstream_socket_.close();
         }
      }

      socket_type downstream_socket_;
      socket_type upstream_socket_;

      enum { max_data_length = 8192 }; //8KB
      unsigned char downstream_data_[max_data_length];
      unsigned char upstream_data_  [max_data_length];

      boost::mutex mutex_;
   public:

      class acceptor
      {
      public:

         acceptor(boost::asio::io_service& io_service,
                  const string& local_host, unsigned short local_port,
                  const string& upstream_host, unsigned short upstream_port)
         : io_service_(io_service),
           localhost_address(boost::asio::ip::address_v4::from_string(local_host)),
           acceptor_(io_service_,ip::tcp::endpoint(localhost_address,local_port)),
           upstream_port_(upstream_port),
           upstream_host_(upstream_host)
         {}

         bool accept_connections()
         {
            try
            {
               session_ = boost::shared_ptr<bridge>(new bridge(io_service_));

               acceptor_.async_accept(session_->downstream_socket(),
                    boost::bind(&acceptor::handle_accept,
                         this,
                         boost::asio::placeholders::error));
            }
            catch(exception& e)
            {
               cerr << "acceptor exception: " << e.what() << endl;
               return false;
            }

            return true;
         }

      private:

         void handle_accept(const boost::system::error_code& error)
         {
            if (!error)
            {
               session_->start(upstream_host_,upstream_port_);

               if (!accept_connections())
               {
                  cerr << "Failure during call to accept." << endl;
               }
            }
            else
            {
               cerr << "Error: " << error.message() << endl;
            }
         }

         boost::asio::io_service& io_service_;
         ip::address_v4 localhost_address;
         ip::tcp::acceptor acceptor_;
         ptr_type session_;
         unsigned short upstream_port_;
         string upstream_host_;
      };

   };
}

void push_regex(char* arg, bool case_sensitive, vector<pair<string,regex>> &v){
   size_t expr_len = (strlen(arg)-2)/2;
   string hex(arg+2);
   string expr;
   if (!unhexlify(hex, expr)){
      cerr << "Regex " << arg << " was not unhexlified successfully" << endl;
      return;
   }
   try{
      if (case_sensitive){
         regex regex(expr);
         #ifdef DEBUG
         cerr << "Added case sensitive regex " << expr_str << endl;
         #endif
         v.push_back(make_pair(string(arg), regex));
      } else {
         regex regex(expr,regex_constants::icase);
         #ifdef DEBUG
         cerr << "Added case insensitive regex " << expr_str << endl;
         #endif
         v.push_back(make_pair(string(arg), regex));
      }
   } catch(...){
      cerr << "Regex " << expr << " was not compiled successfully" << endl;
   }
}


void update_regex(){
   std::unique_lock<std::mutex> lck(update_mutex);
   fstream fd;
   fd.open(config_file,ios::in); 
   if (!fd.is_open()){
	   cerr << "Error: config file couln't be opened" << endl;
      exit(1);
	}

   regex_rules *regex_new_config = new regex_rules();

   string line;
   while(getline(fd, line)){
		char tp[line.length() +1];
		strcpy(tp, line.c_str());
      if (strlen(tp) >= 2){
         bool case_sensitive = true;
         if(tp[0] == '0'){
            case_sensitive = false;
         }
         switch(tp[1]){
            case 'C': { // Client to server Blacklist
               push_regex(tp, case_sensitive, regex_new_config->regex_c_s_b);
               break;
            }
            case 'c': { // Client to server Whitelist
               push_regex(tp, case_sensitive, regex_new_config->regex_c_s_w);
               break;
            }
            case 'S': { // Server to client Blacklist
               push_regex(tp, case_sensitive, regex_new_config->regex_s_c_b);
               break;
            }
            case 's': { // Server to client Whitelist
               push_regex(tp, case_sensitive, regex_new_config->regex_s_c_w);
               break;
            }
         }
      }
   }

   regex_config.reset(regex_new_config);
}

void signal_handler(int signal_num)
{
   if (signal_num == SIGUSR1){
      #ifdef DEBUG
      cerr << "Updating configurtation" << endl;
      #endif
      update_regex();
   }else if(signal_num == SIGTERM){
      if (ios_loop != nullptr) ios_loop->stop();
      #ifdef DEBUG
      cerr << "Close Requested" << endl;
      #endif
      exit(0);
   }
}

int main(int argc, char* argv[])
{
   if (argc < 6)
   {
      cerr << "usage: tcpproxy_server <local host ip> <local port> <forward host ip> <forward port> <config_file>" << endl;
      return 1;
   }

   const unsigned short local_port   = static_cast<unsigned short>(::atoi(argv[2]));
   const unsigned short forward_port = static_cast<unsigned short>(::atoi(argv[4]));
   const string local_host      = argv[1];
   const string forward_host    = argv[3];
   
   config_file = argv[5];

   update_regex();
   signal(SIGUSR1, signal_handler);
   signal(SIGTERM, signal_handler);

   boost::asio::io_service ios;
   ios_loop = &ios;

   try
   {
      tcp_proxy::bridge::acceptor acceptor(ios,
                                           local_host, local_port,
                                           forward_host, forward_port);

      acceptor.accept_connections();
      #ifdef MULTI_THREAD
      boost::thread_group tg;
      #ifdef THREAD_NUM
      for (unsigned i = 0; i < THREAD_NUM; ++i)
      #else
      for (unsigned i = 0; i < thread::hardware_concurrency(); ++i)
      #endif
         tg.create_thread(boost::bind(&boost::asio::io_service::run, &ios));

      tg.join_all();
      #else
      ios.run();
      #endif
   }
   catch(exception& e)
   {
      cerr << "Error: " << e.what() << endl;
      return 1;
   }

   return 0;
}

/*
 * [Note] On posix systems the tcp proxy server build command is as follows:
 * c++ -pedantic -ansi -Wall -Werror -O3 -o tcpproxy_server tcpproxy_server.cpp -L/usr/lib -lstdc++ -lpthread -lboost_thread -lboost_system
 */