Add files via upload

pproxy/__init__.py

@@ -1,8 +1,8 @@
-import argparse, time, re, pickle, asyncio, functools, types, os, urllib.parse
+import argparse, time, re, asyncio, functools, types, urllib.parse
 from pproxy import proto
 
 __title__ = 'pproxy'
-__version__ = "1.0.0"
+__version__ = "1.2.0"
 __description__ = "Proxy server that can tunnel among remote servers by regex rules."
 __author__ = "Qian Wenjie"
 __license__ = "MIT License"
@@ -15,12 +15,35 @@
 asyncio.StreamReader.read_n = lambda self, n: asyncio.wait_for(self.readexactly(n), timeout=SOCKET_TIMEOUT)
 asyncio.StreamReader.read_until = lambda self, s: asyncio.wait_for(self.readuntil(s), timeout=SOCKET_TIMEOUT)
 
-async def proxy_handler(reader, writer, protos, auth, rserver, block, auth_tables, cipher, httpget, unix_path, verbose=DUMMY, modstat=lambda r,h:lambda i:DUMMY, **kwargs):
+if not hasattr(asyncio.StreamReader, 'readuntil'): # Python 3.4 and below
+    @asyncio.coroutine
+    def readuntil(self, separator):
+        seplen = len(separator)
+        offset = 0
+        while True:
+            buflen = len(self._buffer)
+            if buflen - offset >= seplen:
+                isep = self._buffer.find(separator, offset)
+                if isep != -1:
+                    break
+                offset = buflen + 1 - seplen
+            if self._eof:
+                chunk = bytes(self._buffer)
+                self._buffer.clear()
+                raise asyncio.IncompleteReadError(chunk, None)
+            yield from self._wait_for_data('readuntil')
+        chunk = self._buffer[:isep + seplen]
+        del self._buffer[:isep + seplen]
+        self._maybe_resume_transport()
+        return bytes(chunk)
+    asyncio.StreamReader.readuntil = readuntil
+
+def proxy_handler(reader, writer, protos, auth, rserver, block, auth_tables, cipher, httpget, unix_path, verbose=DUMMY, modstat=lambda r,h:lambda i:DUMMY, **kwargs):
     try:
         remote_ip = writer.get_extra_info('peername')[0] if not unix_path else None
         reader_cipher = cipher(reader, writer)[0] if cipher else None
-        header = await reader.read_n(1)
-        lproto, host_name, port, initbuf = await proto.parse(protos, reader=reader, writer=writer, header=header, auth=auth, auth_tables=auth_tables, remote_ip=remote_ip, httpget=httpget, reader_cipher=reader_cipher)
+        header = yield from reader.read_n(1)
+        lproto, host_name, port, initbuf = yield from proto.parse(protos, reader=reader, writer=writer, header=header, auth=auth, auth_tables=auth_tables, remote_ip=remote_ip, httpget=httpget, reader_cipher=reader_cipher)
         if host_name is None:
             writer.close()
             return
@@ -33,36 +56,36 @@ async def proxy_handler(reader, writer, protos, auth, rserver, block, auth_table
                 break
         viaproxy = bool(roption)
         if viaproxy:
-            verbose(f'{lproto.__name__} {host_name}:{port} -> {roption.protos[0].__name__} {roption.bind}')
+            verbose('{l.__name__} {}:{} -> {r.protos[0].__name__} {r.bind}'.format(host_name, port, l=lproto, r=roption))
             connect = roption.connect
         else:
-            verbose(f'{lproto.__name__} {host_name}:{port}')
+            verbose('{l.__name__} {}:{}'.format(host_name, port, l=lproto))
             connect = functools.partial(asyncio.open_connection, host=host_name, port=port)
         try:
-            reader_remote, writer_remote = await asyncio.wait_for(connect(), timeout=SOCKET_TIMEOUT)
+            reader_remote, writer_remote = yield from asyncio.wait_for(connect(), timeout=SOCKET_TIMEOUT)
         except asyncio.TimeoutError:
-            raise Exception(f'Connection timeout {rserver}')
+            raise Exception('Connection timeout {}'.format(rserver))
         try:
             if viaproxy:
                 writer_cipher_r = roption.cipher(reader_remote, writer_remote)[1] if roption.cipher else None
-                await roption.protos[0].connect(reader_remote=reader_remote, writer_remote=writer_remote, rauth=roption.auth, host_name=host_name, port=port, initbuf=initbuf, writer_cipher_r=writer_cipher_r)
+                yield from roption.protos[0].connect(reader_remote=reader_remote, writer_remote=writer_remote, rauth=roption.auth, host_name=host_name, port=port, initbuf=initbuf, writer_cipher_r=writer_cipher_r)
             else:
                 writer_remote.write(initbuf)
         except Exception:
             writer_remote.close()
             raise Exception('Unknown remote protocol')
         m = modstat(remote_ip, host_name)
-        asyncio.ensure_future(proto.base.channel(reader_remote, writer, m(2+viaproxy), m(4+viaproxy)))
-        asyncio.ensure_future(lproto.channel(reader, writer_remote, m(viaproxy), DUMMY))
+        asyncio.async(proto.base.channel(reader_remote, writer, m(2+viaproxy), m(4+viaproxy)))
+        asyncio.async(lproto.channel(reader, writer_remote, m(viaproxy), DUMMY))
     except Exception as ex:
         if not isinstance(ex, asyncio.TimeoutError):
-            verbose(f'{str(ex) or "Unsupported protocol"} from {remote_ip}')
+            verbose('{} from {}'.format(str(ex) or "Unsupported protocol", remote_ip))
         try: writer.close()
         except Exception: pass
 
 def pattern_compile(filename):
     with open(filename) as f:
-        return re.compile('|'.join(i.strip() for i in f if i.strip() and not i.startswith('#'))).fullmatch
+        return re.compile('(:?'+''.join('|'.join(i.strip() for i in f if i.strip() and not i.startswith('#')))+')$').match
 
 def uri_compile(uri):
     url = urllib.parse.urlparse(uri)
@@ -102,20 +125,16 @@ def main():
     parser.add_argument('--ssl', dest='sslfile', help='certfile[,keyfile] if server listen in ssl mode')
     parser.add_argument('--pac', dest='pac', help='http PAC path')
     parser.add_argument('--get', dest='gets', default=[], action='append', help='http custom path/file')
-    parser.add_argument('--version', action='version', version=f'%(prog)s {__version__}')
+    parser.add_argument('--version', action='version', version='%(prog)s {}'.format(__version__))
     args = parser.parse_args()
     if not args.listen:
         args.listen.append(uri_compile('http+socks://:8080/'))
-    if os.path.exists('.auth_tables'):
-        with open('.auth_tables', 'rb') as f:
-            args.auth_tables = pickle.load(f)
-    else:
-        args.auth_tables = {}
+    args.auth_tables = {}
     args.httpget = {}
     if args.pac:
-        pactext = 'function FindProxyForURL(u,h){' + (f'var b=/^(:?{args.block.__self__.pattern})$/i;if(b.test(h))return "";' if args.block else '')
+        pactext = 'function FindProxyForURL(u,h){' + ('var b=/^(:?{})$/i;if(b.test(h))return "";'.format(args.block.__self__.pattern) if args.block else '')
         for i, option in enumerate(args.rserver):
-            pactext += (f'var m{i}=/^(:?{option.match.__self__.pattern})$/i;if(m{i}.test(h))' if option.match else '') + f'return "PROXY %(host)s";'
+            pactext += ('var m{1}=/^(:?{0})$/i;if(m{1}.test(h))'.format(option.match.__self__.pattern, i) if option.match else '') + 'return "PROXY %(host)s";'
         args.httpget[args.pac] = pactext+'return "DIRECT";}'
         args.httpget[args.pac+'/all'] = 'function FindProxyForURL(u,h){return "PROXY %(host)s";}'
         args.httpget[args.pac+'/none'] = 'function FindProxyForURL(u,h){return "DIRECT";}'
@@ -130,29 +149,26 @@ def main():
                 option.sslclient.load_cert_chain(*sslfile)
                 option.sslserver.load_cert_chain(*sslfile)
     elif any(map(lambda o: o.sslclient, args.listen)):
-        print(f'You must specify --ssl to listen in ssl mode')
+        print('You must specify --ssl to listen in ssl mode')
         return
     loop = asyncio.get_event_loop()
     if args.v:
         from pproxy import verbose
         verbose.setup(loop, args)
     servers = []
     for option in args.listen:
-        print(f'Serving on {option.bind} by {",".join(i.__name__ for i in option.protos)}', '(SSL)' if option.sslclient else '')
-        handler = functools.partial(proxy_handler, **vars(args), **vars(option))
+        print('Serving on', option.bind, 'by', ",".join(i.__name__ for i in option.protos) + ('(SSL)' if option.sslclient else ''), '({})'.format(option.cipher.name) if option.cipher else '')
+        handler = functools.partial(functools.partial(proxy_handler, **vars(args)), **vars(option))
         try:
             server = loop.run_until_complete(option.server(handler))
             servers.append(server)
         except Exception as ex:
-            print(f'Start server failed.\n\t==> {ex}')
+            print('Start server failed.\n\t==>', ex)
     if servers:
         try:
             loop.run_forever()
         except KeyboardInterrupt:
             print('exit')
-    if args.auth_tables:
-        with open('.auth_tables', 'wb') as f:
-            pickle.dump(args.auth_tables, f, pickle.HIGHEST_PROTOCOL)
     for task in asyncio.Task.all_tasks():
         task.cancel()
     for server in servers:

pproxy/cipher.py

@@ -1,7 +1,7 @@
 import os, hashlib, argparse, hmac
 
 class BaseCipher(object):
-    LIBRARY = True
+    PYTHON = False
     CACHE = {}
     def __init__(self, key, ota=False):
         if self.KEY_LENGTH > 0:
@@ -24,14 +24,15 @@ def encrypt(self, s):
         return self.cipher.encrypt(s)
     def patch_ota_reader(self, reader):
         chunk_id = 0
-        async def patched_read():
+        @asyncio.coroutine
+        def patched_read():
             nonlocal chunk_id
             try:
-                data_len = int.from_bytes(await reader.readexactly(2), 'big')
+                data_len = int.from_bytes((yield from reader.readexactly(2)), 'big')
             except Exception:
                 return None
-            checksum = await reader.readexactly(10)
-            data = await reader.readexactly(data_len)
+            checksum = yield from reader.readexactly(10)
+            data = yield from reader.readexactly(data_len)
             checksum_server = hmac.new(self.iv+chunk_id.to_bytes(4, 'big'), data, 'sha1').digest()
             assert checksum_server[:10] == checksum
             chunk_id += 1
@@ -47,6 +48,9 @@ def patched_write(data):
             chunk_id += 1
             return write(len(data).to_bytes(2, 'big') + checksum[:10] + data)
         writer.write = patched_write
+    @classmethod
+    def name(cls):
+        return cls.__name__.replace('_Cipher', '').replace('_', '-').lower()
 
 class RC4_Cipher(BaseCipher):
     KEY_LENGTH = 16
@@ -98,6 +102,19 @@ class AES_192_CFB8_Cipher(AES_256_CFB8_Cipher):
 class AES_128_CFB8_Cipher(AES_256_CFB8_Cipher):
     KEY_LENGTH = 16
 
+class AES_256_OFB_Cipher(BaseCipher):
+    KEY_LENGTH = 32
+    IV_LENGTH = 16
+    def setup(self):
+        from Crypto.Cipher import AES
+        self.cipher = AES.new(self.key, AES.MODE_OFB, iv=self.iv)
+
+class AES_192_OFB_Cipher(AES_256_OFB_Cipher):
+    KEY_LENGTH = 24
+
+class AES_128_OFB_Cipher(AES_256_OFB_Cipher):
+    KEY_LENGTH = 16
+
 class BF_CFB_Cipher(BaseCipher):
     KEY_LENGTH = 16
     IV_LENGTH = 8
@@ -119,27 +136,27 @@ def setup(self):
         from Crypto.Cipher import DES
         self.cipher = DES.new(self.key, DES.MODE_CFB, iv=self.iv, segment_size=64)
 
-MAP = {name[:-7].replace('_', '-').lower(): cls for name, cls in globals().items() if name.endswith('_Cipher')}
+MAP = {cls.name(): cls for name, cls in globals().items() if name.endswith('_Cipher')}
 
 def get_cipher(cipher_key):
-    from pproxy.cipherpy import MAP as MAP2
-    CIPHER_MAP = dict(list(MAP.items())+list(MAP2.items()))
+    from pproxy.cipherpy import MAP as MAP_PY
     cipher, _, key = cipher_key.partition(':')
     cipher_name, ota, _ = cipher.partition('!')
     if not key:
         raise argparse.ArgumentTypeError('empty key')
-    if cipher_name not in CIPHER_MAP:
-        raise argparse.ArgumentTypeError(f'existing ciphers: {list(sorted(CIPHER_MAP.keys()))}')
-    cipher, key, ota = CIPHER_MAP[cipher_name], key.encode(), bool(ota) if ota else False
-    if cipher.LIBRARY:
+    if cipher_name not in MAP and cipher_name not in MAP_PY:
+        raise argparse.ArgumentTypeError('existing ciphers: {}'.format(sorted(set(MAP)|set(MAP_PY))))
+    key, ota = key.encode(), bool(ota) if ota else False
+    cipher = MAP.get(cipher_name)
+    if cipher:
         try:
             assert __import__('Crypto').version_info >= (3, 4)
         except Exception:
-            if cipher_name+'-py' in CIPHER_MAP:
-                cipher = CIPHER_MAP[cipher_name+'-py']
-                print(f'Switch to python cipher [{cipher_name}-py]')
-            else:
-                raise argparse.ArgumentTypeError(f'this cipher needs library: "pip3 install pycryptodome"')
+            cipher = None
+    if cipher is None:
+        cipher = MAP_PY.get(cipher_name)
+    if cipher is None:
+        raise argparse.ArgumentTypeError('this cipher needs library: "pip3 install pycryptodome"')
     def apply_cipher(reader, writer):
         reader_cipher, writer_cipher = cipher(key, ota=ota), cipher(key, ota=ota)
         reader_cipher._buffer = b''
@@ -162,7 +179,11 @@ def write(s, o=writer.write):
             return o(writer_cipher.encrypt(s))
         reader.feed_data = feed_data
         writer.write = write
+        if reader._buffer:
+            reader._buffer, buf = bytearray(), reader._buffer
+            feed_data(buf)
         return reader_cipher, writer_cipher
+    apply_cipher.name = cipher_name + ('-py' if cipher.PYTHON else '')
     apply_cipher.ota = ota
     return apply_cipher