seg/trabalho5/proxy.py

import socket
import select
from threading import Thread
import time
import sys
import logging
import logging.handlers
from logging import config

CONNECTION_TIMEOUT_SEC = 10
BUFFER_SIZE = 32 * 1024
CURRENT_THREADS = 0
MAX_THREADS = 50
BACKLOG = 50
LOGGING = {
    'version': 1,
    'disable_existing_loggers': False,
    'formatters': {
        'verbose': {
            'format': '%(levelname)s %(message)s'
        },
    },
    'handlers': {
        'stdout': {
            'class': 'logging.StreamHandler',
            'stream': sys.stdout,
            'formatter': 'verbose',
        },
        'sys-logger6': {
            'class': 'logging.handlers.SysLogHandler',
            "address": ("127.0.0.1", 514),
            'facility': "local6",
            'formatter': 'verbose',
        },
    },
    'loggers': {
        'PythonProxy': {
            'handlers': ['sys-logger6', 'stdout'],
            'level': logging.DEBUG,
            'propagate': True,
        },
    }
}
config.dictConfig(LOGGING)

class Logger:
    _instance = None
    def __init__(self):
        if self._instance is not None:
            raise Exception("Logger is a singleton!")

        self.logger = logging.getLogger('PythonProxy')
        self.logger.info('Initiating Proxy logger!')

    def info(self, message:str):
        self.logger.info(message)

    def critical(self, message:str):
        self.logger.critical(message)

    @classmethod
    def instance(self):
        if self._instance is None:
            self._instance = self()
        return self._instance


class Server:
    def __init__(self, host:str, port:int):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.bind((host, port))
        self.sock.listen(BACKLOG)
        print(f"Proxy executando em: http://{host}:{port}")

    def thread_check(self):
        global CURRENT_THREADS, MAX_THREADS
        while True:
            if CURRENT_THREADS >= MAX_THREADS:
                time.sleep(1)
            else:
                return

    def start(self):
        global CURRENT_THREADS
        while True:
            conn, client_addr = self.sock.accept()
            self.thread_check()

            CURRENT_THREADS += 1
            client_connection = ClientConnection(conn, client_addr)
            thread = Thread(target = client_connection.handle_connection)
            CURRENT_THREADS -= 1
            thread.start()

    def __del__(self):
        self.sock.close()

class ClientConnection:
    def __init__(self, client_socket, client_address):
        self.client_socket = client_socket
        self.client_addr = client_address[0]
        self.client_port = client_address[1]

        self.server_socket = None

    def check_monitorando(self, request):
        if b"monitorando" in request:
            body = open('forbidden.html', 'r').read()
            self.client_socket.sendall(b"HTTP/1.1 403 Forbidden\r\n\r\n")
            self.client_socket.sendall(body.encode())
            self.client_socket.close()
            return True
        
        return False

    def get_host_port(self, request):
        request_url = request.split(' ')[1]
        request_host = ""
        request_port = 443 if request_url.startswith('https') else 80

        # Remove protocolo do request
        if request_url.startswith('http'):
            request_host = request_url.split('//')[1]
        else:
            request_host = request_url.split('/')[0]

        # Remove 'www' do request
        if request_host.startswith('www'):
            request_host = request_host.split('www.')[1]

        # Remove porta do request e verifica se é um host com porta
        if ':' in request_host:
            request_port = int(request_host.split(':')[1])
            request_host = request_host.split(':')[0]

        request_host = request_host.split('/')[0]

        return request_host, request_port

    def handle_connection(self):
        request = self.client_socket.recv(BUFFER_SIZE)
        logger = Logger.instance()

        # Verifica se o request é vazio, se for, fecha o socket
        if len(request) == 0:
            self.client_socket.close()
            return

        try:
            # Tenta decodificar o request, se não conseguir, fecha o socket
            raw_request = request.decode()
        except UnicodeDecodeError:
            self.client_socket.close()
            return

        # Verifica se é um request CONNECT
        # Se for, retorna 200 Connection Established, e espera o request novamente
        if "CONNECT" in raw_request:
            self.client_socket.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n")
            request = self.client_socket.recv(BUFFER_SIZE)

        # Retorna host e porta do request
        request_host, request_port = self.get_host_port(raw_request)

        # Verifica se 'monitorando' está no request
        if self.check_monitorando(request):
            logger.info(f"REQUEST [{self.client_addr}:{self.client_port}] to [{request_host}:{request_port}] - 'Monitorando' - 403 Forbidden")
            return

        # Cria socket para o servidor
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.connect((request_host, request_port))
        self.server_socket.send(request)

        '''
        Enquanto houver dados para serem lidos, lê do socket do cliente e envia para o servidor
        e lê do socket do servidor e envia para o cliente

        Se não houver dados para serem lidos, fecha os sockets

        Exemplo de fluxo:
            proxy -> server
            server -> proxy
            proxy -> client
            proxy -> client
            proxy -> client
            server -> proxy
        '''
        data = None
        while True:
            triple = select.select([self.client_socket, self.server_socket], [], [], CONNECTION_TIMEOUT_SEC)[0]
            if not len(triple):
                break
            try:
                if self.server_socket in triple:
                    data = self.server_socket.recv(BUFFER_SIZE)
                    if not data:
                        break

                    try:
                        status_code = data.decode().split('\r\n')[0].split(' ')[1:]
                        status_code = ' '.join(status_code)
                        if is_valid_status_code(status_code):
                            logger.info(f"REQUEST [{self.client_addr}:{self.client_port}] to [{request_host}:{request_port}] - {status_code}")
                    except UnicodeDecodeError:
                        logger.info(f"REQUEST [{self.client_addr}:{self.client_port}] to [{request_host}:{request_port}]")
                        pass

                    self.client_socket.send(data)
                if self.client_socket in triple:
                    data = self.client_socket.recv(BUFFER_SIZE)
                    if not data:
                        break

                    self.server_socket.send(data)
            except ConnectionAbortedError:
                break

        self.server_socket.close()
        self.client_socket.close()

def is_valid_status_code(status_code: str):
    valid_starts = [str(i) for i in range(5)]
    if status_code.startswith(tuple(valid_starts)):
        return True
    return False

def check_code():
    import hashlib

    content = None
    shas256_hash = None
    logger = Logger.instance()

    with open(sys.argv[0], 'rb') as f:
        content = f.read()

    with open('proxy.sha256sum', 'r') as f:
        shas256_hash = f.read()

    hash = hashlib.sha256(content).hexdigest()

    if hash != shas256_hash:
        logger.critical('ERRO: Código do proxy está diferente!')
        exit(1)

    logger.info('Arquivo do proxy verificado!')

if __name__ == '__main__':
    check_code()

    try:
        ser = Server(host="0.0.0.0", port=8080)
        ser.start()
    except KeyboardInterrupt:
        print("Desligando proxy...")
        exit(0)