<commit_before>#! /usr/bin/env python
import socketserver
from config import SERVER_IP, SERVER_PORT
from encrypt import encrypt, decrypt
from select import select
import logging
import socket
import struct

D = {'C': 0}


def safe_recv(sock, size):
    data = sock.recv(size)
    if not data:
        return None

    size_left = size - len(data)
    while size_left > 0:
        time.sleep(0.01)
        _data = sock.recv(size_left)
        if not _data:
            return None
        data += _data
        size_left = size_left - len(_data)
    return data


class LightHandler(socketserver.BaseRequestHandler):
    def handle_tcp(self, remote):
        sock = self.request
        sock_list = [sock, remote]
        while 1:
            read_list, _, _ = select(sock_list, [], [])
            if remote in read_list:
                data = remote.recv(8192)
                if not data:
                    break
                if (sock.send(encrypt(data)) <= 0):
                    logging.error('send to client error')
                    break

            if sock in read_list:
                data = sock.recv(8192)
                if not data:
                    break
                if (remote.send(decrypt(data)) <= 0):
                    logging.info('send to server error')
                    break

    def handle(self):
        data = self.request.recv(1)
        if not data:
            return
        logging.info('client connected [{}]'.format(D['C']))

        addr_length = data[0]
        data_addr = safe_recv(self.request, addr_length)
        if data_addr is None:
            return

        addr = decrypt(data_addr).decode()
        data_port = safe_recv(self.request, 2)
        if data_port is None:
            return
        port = int.from_bytes(data_port, 'big')

        D['C'] += 1
        logging.info("connecting to {}:{}".format(addr, port))
        remote = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            remote.connect((addr, port))
        except:
            logging.error('cannot connect to {}:{}'.format(addr, port))
            return
        logging.info("waiting for {}:{}".format(addr, port))
        try:
            self.handle_tcp(remote)
        finally:
            remote.close()
        D['C'] -= 1
        logging.info('client closed [{}]'.format(D['C']))


class ThreadedTCPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
    allow_reuse_address = True


if __name__ == "__main__":
    logging.basicConfig(
        format='[%(asctime)s] %(levelname)s: %(message)s',
        level=logging.INFO,
        filename='/tmp/lightsocks-server.log',
    )
    server = ThreadedTCPServer((SERVER_IP, SERVER_PORT), LightHandler)
    logging.info('Proxy running at {}:{} ...'.format(
        SERVER_IP, SERVER_PORT))
    server.serve_forever()<commit_msg>tiny fix<commit_after>#! /usr/bin/env python
import socketserver
from config import SERVER_IP, SERVER_PORT
from encrypt import encrypt, decrypt
from select import select
import logging
import socket
import struct

D = {'C': 0}


def safe_recv(sock, size):
    data = sock.recv(size)
    if not data:
        return None

    size_left = size - len(data)
    while size_left > 0:
        time.sleep(0.01)
        _data = sock.recv(size_left)
        if not _data:
            return None
        data += _data
        size_left = size_left - len(_data)
    return data


class LightHandler(socketserver.BaseRequestHandler):
    def handle_tcp(self, remote):
        sock = self.request
        sock_list = [sock, remote]
        while 1:
            read_list, _, _ = select(sock_list, [], [])
            if remote in read_list:
                data = remote.recv(8192)
                if not data:
                    break
                if (sock.send(encrypt(data)) <= 0):
                    logging.error('send to client error')
                    break

            if sock in read_list:
                data = sock.recv(8192)
                if not data:
                    break
                if (remote.send(decrypt(data)) <= 0):
                    logging.info('send to server error')
                    break

    def handle(self):
        data = self.request.recv(1)
        if not data:
            return
        logging.info('client connected [{}]'.format(D['C']))

        addr_length = data[0]
        data_addr = safe_recv(self.request, addr_length)
        if data_addr is None:
            return

        addr = decrypt(data_addr).decode()
        data_port = safe_recv(self.request, 2)
        if data_port is None:
            return
        port = int.from_bytes(data_port, 'big')
        logging.info("connecting to {}:{}".format(addr, port))
        remote = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            remote.connect((addr, port))
        except:
            logging.error('cannot connect to {}:{}'.format(addr, port))
            return

        D['C'] += 1
        logging.info("waiting for {}:{}".format(addr, port))
        try:
            self.handle_tcp(remote)
        finally:
            remote.close()
        D['C'] -= 1
        logging.info('client closed [{}]'.format(D['C']))


class ThreadedTCPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
    allow_reuse_address = True


if __name__ == "__main__":
    logging.basicConfig(
        format='[%(asctime)s] %(levelname)s: %(message)s',
        level=logging.INFO,
        filename='/tmp/lightsocks-server.log',
    )
    server = ThreadedTCPServer((SERVER_IP, SERVER_PORT), LightHandler)
    logging.info('Proxy running at {}:{} ...'.format(
        SERVER_IP, SERVER_PORT))
    server.serve_forever()
{'fraction_non_alphanumeric': 0.11553994732221247, 'fraction_numerical': 0.007374890254609306, 'mean_word_length': 1.9090909090909092, 'pattern_counts': {'":': 2, '<': 7, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3125100', 'n_tokens_mistral': 1714, 'n_tokens_neox': 1643, 'n_words': 400}