123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490 |
- #!/usr/bin/python3
- #
- import socket
- import struct
- import logging
- import time
- import enum
- import select
- import os
- #
- class ProtocolException(Exception):
- pass
- #
- class ProtocolHelper():
- def __init__(self):
- self._buffer = b''
- #
- def set_buffer(self, data):
- self._buffer = data
- #
- def get_buffer(self):
- return self._buffer
- #
- def recv(self, socket, num_bytes):
- data = socket.recv(num_bytes-len(self._buffer))
- self._buffer += data
- if len(self._buffer) == num_bytes:
- return True
- #
- return False
- #
- def send(self, socket):
- n = socket.send(self._buffer)
- self._buffer = self._buffer[n:]
- if len(self._buffer) == 0:
- return True
- #
- return False
- #
- #
- class Protocol():
- def _run_iteration(self, block=True):
- pass
- #
- def run(self, block=True):
- while True:
- finished = self._run_iteration(block=block)
- #
- if finished:
- # protocol is done
- return True
- elif not block:
- # not done the protocol yet, but don't block
- return False
- #
- #
- #
- #
- class Socks4Protocol(Protocol):
- def __init__(self, socket, addr_port, username=None):
- self.socket = socket
- self.addr_port = addr_port
- self.username = username
- #
- self.states = enum.Enum('SOCKS_4_STATES', 'READY_TO_BEGIN CONNECTING_TO_PROXY WAITING_FOR_PROXY DONE')
- self.state = self.states.READY_TO_BEGIN
- #
- self.protocol_helper = None
- #
- def _run_iteration(self, block=True):
- if self.state is self.states.READY_TO_BEGIN:
- self.protocol_helper = ProtocolHelper()
- self.protocol_helper.set_buffer(self.socks_cmd(self.addr_port, self.username))
- self.state = self.states.CONNECTING_TO_PROXY
- #
- if self.state is self.states.CONNECTING_TO_PROXY:
- if self.protocol_helper.send(self.socket):
- self.protocol_helper = ProtocolHelper()
- self.state = self.states.WAITING_FOR_PROXY
- #logging.debug('Waiting for reply from proxy')
- #
- #
- if self.state is self.states.WAITING_FOR_PROXY:
- response_size = 8
- if self.protocol_helper.recv(self.socket, response_size):
- response = self.protocol_helper.get_buffer()
- if response[1] != 0x5a:
- raise ProtocolException('Could not connect to SOCKS proxy, msg: %x'%(response[1],))
- #
- self.state = self.states.DONE
- return True
- #
- #
- return False
- #
- def socks_cmd(self, addr_port, username=None):
- socks_version = 4
- command = 1
- dnsname = b''
- host, port = addr_port
- #
- try:
- username = bytes(username, 'utf8')
- except TypeError:
- pass
- #
- if username is None:
- username = b''
- elif b'\x00' in username:
- raise ProtocolException('Username cannot contain a NUL character.')
- #
- username = username+b'\x00'
- #
- try:
- addr = socket.inet_aton(host)
- except socket.error:
- addr = b'\x00\x00\x00\x01'
- dnsname = bytes(host, 'utf8')+b'\x00'
- #
- return struct.pack('!BBH', socks_version, command, port) + addr + username + dnsname
- #
- #
- class PushDataProtocol(Protocol):
- def __init__(self, socket, total_bytes, data_generator=None, send_max_bytes=1024*512):
- if data_generator is None:
- data_generator = self._default_data_generator
- #
- self.socket = socket
- self.data_generator = data_generator
- self.total_bytes = total_bytes
- self.send_max_bytes = send_max_bytes
- #
- self.states = enum.Enum('PUSH_DATA_STATES', 'READY_TO_BEGIN SEND_INFO PUSH_DATA RECV_CONFIRMATION DONE')
- self.state = self.states.READY_TO_BEGIN
- #
- self.bytes_written = 0
- self.protocol_helper = None
- #
- def _run_iteration(self, block=True):
- if self.state is self.states.READY_TO_BEGIN:
- info = self.total_bytes.to_bytes(8, byteorder='big', signed=False)
- info += self.send_max_bytes.to_bytes(8, byteorder='big', signed=False)
- self.protocol_helper = ProtocolHelper()
- self.protocol_helper.set_buffer(info)
- self.state = self.states.SEND_INFO
- #
- if self.state is self.states.SEND_INFO:
- if self.protocol_helper.send(self.socket):
- self.state = self.states.PUSH_DATA
- #
- #
- if self.state is self.states.PUSH_DATA:
- max_block_size = self.send_max_bytes
- bytes_needed = min(max_block_size, self.total_bytes-self.bytes_written)
- data = self.data_generator(self.bytes_written, bytes_needed)
- n = self.socket.send(data)
- self.bytes_written += n
- if self.bytes_written >= self.total_bytes:
- # finished sending the data
- logging.debug('Finished sending the data (%d bytes).', self.bytes_written)
- self.protocol_helper = ProtocolHelper()
- self.state = self.states.RECV_CONFIRMATION
- #
- #
- if self.state is self.states.RECV_CONFIRMATION:
- response_size = 8
- if self.protocol_helper.recv(self.socket, response_size):
- response = self.protocol_helper.get_buffer()
- if response != b'RECEIVED':
- raise ProtocolException('Did not receive the expected message: {}'.format(response))
- #
- self.state = self.states.DONE
- return True
- #
- #
- return False
- #
- def _default_data_generator(self, index, bytes_needed):
- return b'0'*bytes_needed
- #
- #
- class PullDataProtocol(Protocol):
- def __init__(self, socket):
- self.socket = socket
- #
- self.states = enum.Enum('PULL_DATA_STATES', 'READY_TO_BEGIN RECV_INFO PULL_DATA SEND_CONFIRMATION DONE')
- self.state = self.states.READY_TO_BEGIN
- #
- self.data_size = None
- self.recv_max_bytes = None
- self.bytes_read = 0
- self.protocol_helper = None
- #
- def _run_iteration(self, block=True):
- if self.state is self.states.READY_TO_BEGIN:
- self.protocol_helper = ProtocolHelper()
- self.state = self.states.RECV_INFO
- #
- if self.state is self.states.RECV_INFO:
- info_size = 16
- if self.protocol_helper.recv(self.socket, info_size):
- response = self.protocol_helper.get_buffer()
- self.data_size = int.from_bytes(response[0:8], byteorder='big', signed=False)
- self.recv_max_bytes = int.from_bytes(response[8:16], byteorder='big', signed=False)
- self.state = self.states.PULL_DATA
- #
- #
- if self.state is self.states.PULL_DATA:
- max_block_size = self.recv_max_bytes
- bytes_needed = min(max_block_size, self.data_size-self.bytes_read)
- data = self.socket.recv(bytes_needed)
- self.bytes_read += len(data)
- #logging.debug('Read %d bytes', self.bytes_read)
- if self.bytes_read == self.data_size:
- # finished receiving the data
- logging.debug('Finished receiving the data.')
- self.protocol_helper = ProtocolHelper()
- self.protocol_helper.set_buffer(b'RECEIVED')
- self.state = self.states.SEND_CONFIRMATION
- #
- #
- if self.state is self.states.SEND_CONFIRMATION:
- if self.protocol_helper.send(self.socket):
- self.state = self.states.DONE
- return True
- #
- #
- return False
- #
- #
- class PullDataProtocolWithMetrics(PullDataProtocol):
- def __init__(self, *args, **kwargs):
- super().__init__(*args, **kwargs)
- #
- self.time_of_first_byte = None
- self.time_of_last_byte = None
- #
- def _run_iteration(self, *args, **kwargs):
- data = super()._run_iteration(*args, **kwargs)
- #
- if self.bytes_read != 0 and self.time_of_first_byte is None:
- self.time_of_first_byte = time.time()
- #
- if self.bytes_read == self.data_size and self.time_of_last_byte is None:
- self.time_of_last_byte = time.time()
- #
- return data
- #
- def calc_transfer_rate(self):
- """ Returns bytes/s. """
- assert self.data_size is not None and self.time_of_first_byte is not None and self.time_of_last_byte is not None
- return self.data_size/(self.time_of_last_byte-self.time_of_first_byte)
- #
- #
- class SendDataProtocol(Protocol):
- def __init__(self, socket, data):
- self.socket = socket
- self.send_data = data
- #
- self.states = enum.Enum('SEND_DATA_STATES', 'READY_TO_BEGIN SEND_INFO SEND_DATA RECV_CONFIRMATION DONE')
- self.state = self.states.READY_TO_BEGIN
- #
- self.protocol_helper = None
- #
- def _run_iteration(self, block=True):
- if self.state is self.states.READY_TO_BEGIN:
- info_size = 20
- info = len(self.send_data).to_bytes(info_size, byteorder='big', signed=False)
- self.protocol_helper = ProtocolHelper()
- self.protocol_helper.set_buffer(info)
- self.state = self.states.SEND_INFO
- #
- if self.state is self.states.SEND_INFO:
- if self.protocol_helper.send(self.socket):
- self.protocol_helper = ProtocolHelper()
- self.protocol_helper.set_buffer(self.send_data)
- self.state = self.states.SEND_DATA
- #
- #
- if self.state is self.states.SEND_DATA:
- if self.protocol_helper.send(self.socket):
- self.protocol_helper = ProtocolHelper()
- self.state = self.states.RECV_CONFIRMATION
- #
- #
- if self.state is self.states.RECV_CONFIRMATION:
- response_size = 8
- if self.protocol_helper.recv(self.socket, response_size):
- response = self.protocol_helper.get_buffer()
- if response != b'RECEIVED':
- raise ProtocolException('Did not receive the expected message: {}'.format(response))
- #
- self.state = self.states.DONE
- return True
- #
- #
- return False
- #
- #
- class ReceiveDataProtocol(Protocol):
- def __init__(self, socket):
- self.socket = socket
- #
- self.states = enum.Enum('RECV_DATA_STATES', 'READY_TO_BEGIN RECV_INFO RECV_DATA SEND_CONFIRMATION DONE')
- self.state = self.states.READY_TO_BEGIN
- #
- self.protocol_helper = None
- self.data_size = None
- self.received_data = None
- #
- def _run_iteration(self, block=True):
- if self.state is self.states.READY_TO_BEGIN:
- self.protocol_helper = ProtocolHelper()
- self.state = self.states.RECV_INFO
- #
- if self.state is self.states.RECV_INFO:
- info_size = 20
- if self.protocol_helper.recv(self.socket, info_size):
- response = self.protocol_helper.get_buffer()
- self.data_size = int.from_bytes(response, byteorder='big', signed=False)
- self.protocol_helper = ProtocolHelper()
- self.state = self.states.RECV_DATA
- #
- #
- if self.state is self.states.RECV_DATA:
- if self.protocol_helper.recv(self.socket, self.data_size):
- response = self.protocol_helper.get_buffer()
- self.received_data = response
- self.protocol_helper = ProtocolHelper()
- self.protocol_helper.set_buffer(b'RECEIVED')
- self.state = self.states.SEND_CONFIRMATION
- #
- #
- if self.state is self.states.SEND_CONFIRMATION:
- if self.protocol_helper.send(self.socket):
- self.state = self.states.DONE
- return True
- #
- #
- return False
- #
- #
- class ServerListener():
- "A TCP listener, binding, listening and accepting new connections."
- def __init__(self, endpoint, accept_callback):
- self.callback = accept_callback
- #
- self.s = socket.socket()
- self.s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
- self.s.bind(endpoint)
- self.s.listen(0)
- #
- def accept(self):
- newsock, endpoint = self.s.accept()
- logging.debug("New client from %s:%d (fd=%d)",
- endpoint[0], endpoint[1], newsock.fileno())
- self.callback(newsock)
- #
- #
- class SimpleClientConnectionProtocol(Protocol):
- def __init__(self, endpoint, total_bytes, data_generator=None, proxy=None, username=None):
- self.endpoint = endpoint
- self.data_generator = data_generator
- self.total_bytes = total_bytes
- self.proxy = proxy
- self.username = username
- #
- self.states = enum.Enum('CLIENT_CONN_STATES', 'READY_TO_BEGIN CONNECT_TO_PROXY PUSH_DATA DONE')
- self.state = self.states.READY_TO_BEGIN
- #
- self.socket = socket.socket()
- self.sub_protocol = None
- #
- if self.proxy is None:
- logging.debug('Socket %d connecting to endpoint %r...', self.socket.fileno(), self.endpoint)
- self.socket.connect(self.endpoint)
- else:
- logging.debug('Socket %d connecting to proxy %r...', self.socket.fileno(), self.proxy)
- self.socket.connect(self.proxy)
- #
- #
- def _run_iteration(self, block=True):
- if self.state is self.states.READY_TO_BEGIN:
- if self.proxy is None:
- self.sub_protocol = PushDataProtocol(self.socket, self.total_bytes, self.data_generator)
- self.state = self.states.PUSH_DATA
- else:
- self.sub_protocol = Socks4Protocol(self.socket, self.endpoint, username=self.username)
- self.state = self.states.CONNECT_TO_PROXY
- #
- #
- if self.state is self.states.CONNECT_TO_PROXY:
- if self.sub_protocol.run(block=block):
- self.sub_protocol = PushDataProtocol(self.socket, self.total_bytes, self.data_generator)
- self.state = self.states.PUSH_DATA
- #
- #
- if self.state is self.states.PUSH_DATA:
- if self.sub_protocol.run(block=block):
- self.state = self.states.DONE
- return True
- #
- #
- return False
- #
- #
- class SimpleServerConnectionProtocol(Protocol):
- def __init__(self, socket, conn_id, bandwidth_callback=None):
- self.socket = socket
- self.conn_id = conn_id
- self.bandwidth_callback = bandwidth_callback
- #
- self.states = enum.Enum('SERVER_CONN_STATES', 'READY_TO_BEGIN PULL_DATA DONE')
- self.state = self.states.READY_TO_BEGIN
- #
- self.sub_protocol = None
- #
- def _run_iteration(self, block=True):
- if self.state is self.states.READY_TO_BEGIN:
- self.sub_protocol = PullDataProtocolWithMetrics(self.socket)
- self.state = self.states.PULL_DATA
- #
- if self.state is self.states.PULL_DATA:
- if self.sub_protocol.run(block=block):
- self.state = self.states.DONE
- if self.bandwidth_callback:
- self.bandwidth_callback(self.conn_id, self.sub_protocol.data_size, self.sub_protocol.calc_transfer_rate())
- #
- return True
- #
- #
- return False
- #
- #
- if __name__ == '__main__':
- import sys
- logging.basicConfig(level=logging.DEBUG)
- #
- if sys.argv[1] == 'client':
- endpoint = ('127.0.0.1', 4747)
- proxy = ('127.0.0.1', 9003)
- #proxy = None
- username = bytes([x for x in os.urandom(12) if x != 0])
- #username = None
- data_MB = 40
- #
- client = SimpleClientConnectionProtocol(endpoint, data_MB*2**20, proxy=proxy, username=username)
- client.run()
- elif sys.argv[1] == 'server':
- import multiprocessing
- import queue
- #
- endpoint = ('127.0.0.1', 4747)
- processes = []
- conn_counter = [0]
- #
- def bw_callback(conn_id, data_size, transfer_rate):
- logging.info('Avg Transferred (MB): %.4f', data_size/(1024**2))
- logging.info('Avg Transfer rate (MB/s): %.4f', transfer_rate/(1024**2))
- #
- def start_server_conn(socket, conn_id):
- server = SimpleServerConnectionProtocol(socket, conn_id, bandwidth_callback=bw_callback)
- try:
- server.run()
- except KeyboardInterrupt:
- socket.close()
- #
- #
- def accept_callback(socket):
- conn_id = conn_counter[0]
- conn_counter[0] += 1
- #
- p = multiprocessing.Process(target=start_server_conn, args=(socket, conn_id))
- processes.append(p)
- p.start()
- #
- l = ServerListener(endpoint, accept_callback)
- #
- try:
- while True:
- l.accept()
- #
- except KeyboardInterrupt:
- print()
- #
- for p in processes:
- p.join()
- #
- #
- #
|