Merge pull request #175 from vforvideo/admin

feat: httpadmin protocol for external control

Author: qwj

pproxy/admin.py

@@ -0,0 +1,39 @@
+import json
+import asyncio
+config = {}
+
+
+async def reply_http(reply, ver, code, content):
+    await reply(code, f'{ver} {code}\r\nConnection: close\r\nContent-Type: text/plain\r\nCache-Control: max-age=900\r\nContent-Length: {len(content)}\r\n\r\n'.encode(), content, True)
+
+
+async def status_handler(reply, **kwarg):
+    method = kwarg.get('method')
+    if method == 'GET':
+        data = {"status": "ok"}
+        value = json.dumps(data).encode()
+        ver = kwarg.get('ver')
+        await reply_http(reply, ver, '200 OK', value)
+
+
+async def configs_handler(reply, **kwarg):
+    method = kwarg.get('method')
+    ver = kwarg.get('ver')
+
+    if method == 'GET':
+        data = {"argv": config['argv']}
+        value = json.dumps(data).encode()
+        await reply_http(reply, ver, '200 OK', value)
+    elif method == 'POST':
+        config['argv'] = kwarg.get('content').decode().split(' ')
+        config['reload'] = True
+        data = {"result": 'ok'}
+        value = json.dumps(data).encode()
+        await reply_http(reply, ver, '200 OK', value)
+        raise KeyboardInterrupt
+
+
+httpget = {
+    '/status': status_handler,
+    '/configs': configs_handler,
+}

pproxy/proto.py

@@ -1,5 +1,5 @@
 import asyncio, socket, urllib.parse, time, re, base64, hmac, struct, hashlib, io, os
-
+from . import admin
 HTTP_LINE = re.compile('([^ ]+) +(.+?) +(HTTP/[^ ]+)$')
 packstr = lambda s, n=1: len(s).to_bytes(n, 'big') + s
 
@@ -404,6 +404,38 @@ async def connect(self, reader_remote, writer_remote, rauth, host_name, port, my
 class H3(H2):
     pass
 
+
+class HTTPAdmin(HTTP):
+    async def accept(self, reader, user, writer, **kw):
+        lines = await reader.read_until(b'\r\n\r\n')
+        headers = lines[:-4].decode().split('\r\n')
+        method, path, ver = HTTP_LINE.match(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)
+        async def reply(code, message, body=None, wait=False):
+            writer.write(message)
+            if body:
+                writer.write(body)
+            if wait:
+                await writer.drain()
+
+        content_length = int(headers.get('Content-Length','0'))
+        content = ''
+        if content_length > 0:
+            content = await reader.read_n(content_length)
+
+        url = urllib.parse.urlparse(path)
+        if url.hostname is not None:
+            raise Exception(f'HTTP Admin Unsupported hostname')
+        if method in ["GET", "POST", "PUT", "PATCH", "DELETE"]:
+            for path, handler in admin.httpget.items():
+                if path == url.path:
+                    await handler(reply=reply, ver=ver, method=method, headers=headers, lines=lines, content=content)
+                    raise Exception('Connection closed')
+            raise Exception(f'404 {method} {url.path}')
+        raise Exception(f'405 {method} not allowed')
+        
+
 class SSH(BaseProtocol):
     async def connect(self, reader_remote, writer_remote, rauth, host_name, port, myhost, **kw):
         pass
@@ -567,7 +599,7 @@ def udp_accept(protos, data, **kw):
             return (proto,) + ret
     raise Exception(f'Unsupported protocol {data[:10]}')
 
-MAPPINGS = dict(direct=Direct, http=HTTP, httponly=HTTPOnly, ssh=SSH, socks5=Socks5, socks4=Socks4, socks=Socks5, ss=SS, ssr=SSR, redir=Redir, pf=Pf, tunnel=Tunnel, echo=Echo, ws=WS, trojan=Trojan, h2=H2, h3=H3, ssl='', secure='', quic='')
+MAPPINGS = dict(direct=Direct, http=HTTP, httponly=HTTPOnly, httpadmin=HTTPAdmin, ssh=SSH, socks5=Socks5, socks4=Socks4, socks=Socks5, ss=SS, ssr=SSR, redir=Redir, pf=Pf, tunnel=Tunnel, echo=Echo, ws=WS, trojan=Trojan, h2=H2, h3=H3, ssl='', secure='', quic='')
 MAPPINGS['in'] = ''
 
 def get_protos(rawprotos):

pproxy/server.py

@@ -1,5 +1,7 @@
-import argparse, time, re, asyncio, functools, base64, random, urllib.parse, socket
+import argparse, time, re, asyncio, functools, base64, random, urllib.parse, socket, sys
 from . import proto
+from . import admin
+
 from .__doc__ import *
 
 SOCKET_TIMEOUT = 60
@@ -892,6 +894,8 @@ def print_server_started(option, server, print_fn):
         print_fn(option, bind)
 
 def main(args = None):
+    origin_argv = sys.argv[1:] if args is None else args
+
     parser = argparse.ArgumentParser(description=__description__+'\nSupported protocols: http,socks4,socks5,shadowsocks,shadowsocksr,redirect,pf,tunnel', epilog=f'Online help: <{__url__}>')
     parser.add_argument('-l', dest='listen', default=[], action='append', type=proxies_by_uri, help='tcp server uri (default: http+socks4+socks5://:8080/)')
     parser.add_argument('-r', dest='rserver', default=[], action='append', type=proxies_by_uri, help='tcp remote server uri (default: direct)')
@@ -953,6 +957,7 @@ def main(args = None):
         from . import verbose
         verbose.setup(loop, args)
     servers = []
+    admin.config.update({'argv': origin_argv, 'servers': servers, 'args': args, 'loop': loop})
     def print_fn(option, bind=None):
         print('Serving on', (bind or option.bind), 'by', ",".join(i.name for i in option.protos) + ('(SSL)' if option.sslclient else ''), '({}{})'.format(option.cipher.name, ' '+','.join(i.name() for i in option.cipher.plugins) if option.cipher and option.cipher.plugins else '') if option.cipher else '')
     for option in args.listen:
@@ -1004,6 +1009,9 @@ def print_fn(option, bind=None):
         if hasattr(server, 'wait_closed'):
             loop.run_until_complete(server.wait_closed())
     loop.run_until_complete(loop.shutdown_asyncgens())
+    if admin.config.get('reload', False):
+        admin.config['reload'] = False
+        main(admin.config['argv'])
     loop.close()
 
 if __name__ == '__main__':