diff --git a/tsproxy.py b/tsproxy.py index 2e55aeb..cd5b948 100644 --- a/tsproxy.py +++ b/tsproxy.py @@ -18,12 +18,13 @@ import gc import logging import platform + try: - from Queue import Queue - from Queue import Empty + from Queue import Empty, Queue except ImportError: from queue import Queue from queue import Empty + import re import signal import socket @@ -48,7 +49,13 @@ REMOVE_TCP_OVERHEAD = 1460.0 / 1500.0 lock = threading.Lock() background_activity_count = 0 -current_time = time.clock if sys.platform == "win32" else time.time +if sys.platform == "win32": + try: + current_time = time.perf_counter + except AttributeError: + current_time = time.clock +else: + current_time = time.time try: import monotonic current_time = monotonic.monotonic @@ -56,10 +63,27 @@ pass +if sys.version_info.major == 3: + # In Python 2, data from/to the socket are stored in character strings, + # and the built-in ord() and chr() functions are used to convert between + # characters and integers. + # + # In Python 3, data are stored in bytes, and we need to redefine ord and + # chr functions to make it work. + + def ord(x): + # In Python 3, indexing a byte string returns an int, no conversion needed. + return x + + def chr(x): + # Convert a byte into bytes of length 1. + return bytes([x]) + + def PrintMessage(msg): # Print the message to stdout & flush to make sure that the message is not # buffered when tsproxy is run as a subprocess. - sys.stdout.write(msg) + sys.stdout.write(msg + '\n') sys.stdout.flush() ######################################################################################################################## @@ -234,7 +258,7 @@ def __init__(self, client_id): asyncore.dispatcher.__init__(self) self.client_id = client_id self.state = self.STATE_IDLE - self.buffer = '' + self.buffer = b'' self.addr = None self.dns_thread = None self.hostname = None @@ -336,9 +360,9 @@ def HandleResolve(self, message): if 'port' in message: self.port = message['port'] logging.info('[{0:d}] Resolving {1}:{2:d}'.format(self.client_id, self.hostname, self.port)) - if self.hostname == 'localhost': - self.hostname = '127.0.0.1' - if self.hostname == '127.0.0.1': + if self.hostname == b'localhost': + self.hostname = b'127.0.0.1' + if self.hostname == b'127.0.0.1': logging.info('[{0:d}] Connection to localhost detected'.format(self.client_id)) is_localhost = True if (dest_addresses is not None) and (not is_localhost or map_localhost): @@ -427,7 +451,7 @@ def __init__(self, connected_socket, client_id): self.hostname = None self.port = None self.requested_address = None - self.buffer = '' + self.buffer = b'' self.socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 128 * 1024) self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 128 * 1024) @@ -615,6 +639,7 @@ def ProcessCommand(self, input): global REMOVE_TCP_OVERHEAD global port_mappings global server + global must_exit if len(input): ok = False try: @@ -652,6 +677,9 @@ def ProcessCommand(self, input): if command[1].lower() == 'mapports' or command[1].lower() == 'all': port_mappings = {} ok = True + elif command[0].lower() == 'exit': + must_exit = True + ok = True if ok: needs_flush = True @@ -682,7 +710,7 @@ def main(): global REMOVE_TCP_OVERHEAD parser = argparse.ArgumentParser(description='Traffic-shaping socks5 proxy.', prog='tsproxy') - parser.add_argument('-v', '--verbose', action='count', help="Increase verbosity (specify multiple times for more). -vvvv for full debug output.") + parser.add_argument('-v', '--verbose', action='count', default=0, help="Increase verbosity (specify multiple times for more). -vvvv for full debug output.") parser.add_argument('--logfile', help="Write log messages to given file instead of stdout.") parser.add_argument('-b', '--bind', default='localhost', help="Server interface address (defaults to localhost).") parser.add_argument('-p', '--port', type=int, default=1080, help="Server port (defaults to 1080, use 0 for randomly assigned).") @@ -717,7 +745,7 @@ def main(): # Parse any port mappings if options.mapports: SetPortMappings(options.mapports) - + if options.nodnscache: dns_cache = None