Add files via upload

Author: qwj

pproxy/__init__.py

@@ -2,7 +2,7 @@
 from pproxy import proto
 
 __title__ = 'pproxy'
-__version__ = "0.9.4"
+__version__ = "0.9.6"
 __description__ = "Proxy server that can tunnel among remote servers by regex rules."
 __author__ = "Qian Wenjie"
 __license__ = "MIT License"
@@ -15,12 +15,12 @@
 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, pac, pactext, unix_path, verbose=DUMMY, modstat=lambda r,h:lambda i:DUMMY, **kwargs):
+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):
     try:
         remote_ip = writer.get_extra_info('peername')[0] if not unix_path else None
-        reader_cipher = (await cipher(reader, writer))[0] if cipher 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, pac=pac, pactext=pactext, reader_cipher=reader_cipher)
+        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)
         if host_name is None:
             writer.close()
             return
@@ -44,7 +44,7 @@ async def proxy_handler(reader, writer, protos, auth, rserver, block, auth_table
             raise Exception(f'Connection timeout {rserver}')
         try:
             if viaproxy:
-                writer_cipher_r = (await roption.cipher(reader_remote, writer_remote))[1] if roption.cipher else None
+                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)
             else:
                 writer_remote.write(initbuf)
@@ -60,8 +60,8 @@ async def proxy_handler(reader, writer, protos, auth, rserver, block, auth_table
         try: writer.close()
         except Exception: pass
 
-def pattern_compile(file_name):
-    with open(file_name) as f:
+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
 
 def uri_compile(uri):
@@ -100,7 +100,8 @@ def main():
     parser.add_argument('-b', dest='block', type=pattern_compile, help='block regex rules')
     parser.add_argument('-v', dest='v', action='store_true', help='print verbose output')
     parser.add_argument('--ssl', dest='sslfile', help='certfile[,keyfile] if server listen in ssl mode')
-    parser.add_argument('--pac', dest='pac', help='http pac file path')
+    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__}')
     args = parser.parse_args()
     if not args.listen:
@@ -110,13 +111,18 @@ def main():
             args.auth_tables = pickle.load(f)
     else:
         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 '')
         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";'
-        args.pactext = (pactext+'return "DIRECT";}', 'function FindProxyForURL(u,h){{return "PROXY %(host)s";}}', 'function FindProxyForURL(u,h){return "DIRECT";}')
-    else:
-        args.pactext = None
+        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";}'
+    for gets in args.gets:
+        path, filename = gets.split(',', 1)
+        with open(filename, 'r') as f:
+            args.httpget[path] = f.read()
     if args.sslfile:
         sslfile = args.sslfile.split(',')
         for option in args.listen:

pproxy/ciphers.py

@@ -5,7 +5,7 @@
 
 class BaseCipher(object):
     CACHE = {}
-    def __init__(self, key, iv=None, ota=False):
+    def __init__(self, key, ota=False):
         if self.KEY_LENGTH > 0:
             self.key = self.CACHE.get(b'key'+key)
             if self.key is None:
@@ -15,8 +15,10 @@ def __init__(self, key, iv=None, ota=False):
                 self.key = self.CACHE[b'key'+key] = b''.join(keybuf)[:self.KEY_LENGTH]
         else:
             self.key = key
-        self.iv = os.urandom(self.IV_LENGTH) if iv is None else iv
         self.ota = ota
+        self.iv = None
+    def setup_iv(self, iv=None):
+        self.iv = os.urandom(self.IV_LENGTH) if iv is None else iv
         self.setup()
     def decrypt(self, s):
         return self.cipher.decrypt(s)
@@ -143,13 +145,27 @@ def get_cipher(cipher_key):
     if cipher not in MAPPINGS:
         raise argparse.ArgumentTypeError(f'existing ciphers: {list(MAPPINGS.keys())}')
     cipher, key, ota = MAPPINGS[cipher], key.encode(), bool(ota) if ota else False
-    async def apply_cipher(reader, writer):
-        writer_cipher = cipher(key, ota=ota)
-        writer.write(writer_cipher.iv)
-        writer.write = lambda s, o=writer.write, p=writer_cipher.encrypt: o(p(s))
-        reader_cipher = cipher(key, await reader.read_n(len(writer_cipher.iv)), ota=ota)
-        reader._buffer = bytearray(reader_cipher.decrypt(bytes(reader._buffer)))
-        reader.feed_data = lambda s, o=reader.feed_data, p=reader_cipher.decrypt: o(p(s))
+    def apply_cipher(reader, writer):
+        reader_cipher, writer_cipher = cipher(key, ota=ota), cipher(key, ota=ota)
+        def feed_data(s, o=reader.feed_data):
+            if not reader_cipher.iv:
+                s = bytes(reader._buffer + s)
+                if len(s) >= reader_cipher.IV_LENGTH:
+                    reader_cipher.setup_iv(s[:reader_cipher.IV_LENGTH])
+                    s = reader_cipher.decrypt(s[reader_cipher.IV_LENGTH:])
+                    reader._buffer.clear()
+            else:
+                s = reader_cipher.decrypt(s)
+            return o(s)
+        def write(s, o=writer.write):
+            if not s:
+                return
+            if not writer_cipher.iv:
+                writer_cipher.setup_iv()
+                o(writer_cipher.iv)
+            return o(writer_cipher.encrypt(s))
+        reader.feed_data = feed_data
+        writer.write = write
         return reader_cipher, writer_cipher
     apply_cipher.ota = ota
     return apply_cipher

pproxy/proto.py

@@ -115,18 +115,21 @@ def correct_header(header, **kw):
         return header.isalpha()
 
     @staticmethod
-    async def parse(header, reader, writer, auth, auth_tables, remote_ip, pac, pactext, **kw):
+    async def parse(header, reader, writer, auth, auth_tables, remote_ip, httpget, **kw):
         lines = header + await reader.read_until(b'\r\n\r\n')
         headers = lines[:-4].decode().split('\r\n')
         method, path, ver = HTTP_LINE.fullmatch(headers.pop(0)).groups()
         lines = '\r\n'.join(i for i in headers if not i.startswith('Proxy-'))
         headers = dict(i.split(': ', 1) for i in headers if ': ' in i)
-        if method == 'GET' and pac and path.startswith(pac):
-            auth_tables[remote_ip] = time.time()
-            pacidx = 1 if 'all' in path else (2 if 'none' in path else 0)
-            text = (pactext[pacidx] % dict(host=headers["Host"])).encode()
-            writer.write(f'{ver} 200 OK\r\nConnection: close\r\nContent-Type: application/x-ns-proxy-autoconfig\r\nCache-Control: max-age=900\r\nContent-Length: {len(text)}\r\n\r\n'.encode() + text)
-            return None, None, None
+        url = urllib.parse.urlparse(path)
+        if method == 'GET' and not url.hostname:
+            for path, text in httpget.items():
+                if url.path == path:
+                    auth_tables[remote_ip] = time.time()
+                    text = (text % dict(host=headers["Host"])).encode()
+                    writer.write(f'{ver} 200 OK\r\nConnection: close\r\nContent-Type: text/plain\r\nCache-Control: max-age=900\r\nContent-Length: {len(text)}\r\n\r\n'.encode() + text)
+                    return None, None, None
+            raise Exception(f'404 {method} {path}')
         if auth:
             pauth = headers.get('Proxy-Authorization', None)
             httpauth = 'Basic ' + base64.b64encode(auth).decode()
@@ -142,8 +145,6 @@ async def parse(header, reader, writer, auth, auth_tables, remote_ip, pac, pacte
         else:
             url = urllib.parse.urlparse(path)
             host_name = url.hostname
-            if not host_name:
-                raise Exception(f'404 {method} {path}')
             port = url.port or 80
             newpath = url._replace(netloc='', scheme='').geturl()
             return host_name, port, f'{method} {newpath} {ver}\r\n{lines}\r\n\r\n'.encode()