fix pypi dependencies + small howto

Author: Yurzs

README.md

@@ -1,2 +1,13 @@
 # aMTProxy
 Asynchronous Python MTProto proxy server [beta]
+
+# Installation
+To get the latest build use `pip install amtproxy`
+
+# Usage
+
+`from amtproxy.server import start_server`\
+`start_server(port=1234, secret='32-digit hex value')`
+
+# TBD
+- [ ] `dd` secret prefix

amtproxy/config.py

@@ -1,5 +1,3 @@
-telegram_ips = [ "149.154.175.50", "149.154.167.51", "149.154.175.100", "149.154.167.91", "149.154.171.5"]
+telegram_ips = ["149.154.175.50", "149.154.167.51", "149.154.175.100", "149.154.167.91", "149.154.171.5"]
 mtproto_footprint = b'\xef\xef\xef\xef'
 obf_footprint = b'\xdd\xdd\xdd\xdd'
-secret = '1234567890abc1234567890abcef1234'
-

amtproxy/dispatcher.py

@@ -1,12 +1,13 @@
-from amtproxy.encryption.crypto import Encryptor, Decryptor
-from amtproxy.config import secret, telegram_ips, mtproto_footprint, obf_footprint
-from amtproxy.exceptions import WrongServerID
-from amtproxy.protocols.telegram import TelegramProtocol
+from .encryption.crypto import Encryptor, Decryptor
+from .config import telegram_ips, mtproto_footprint, obf_footprint
+from .exceptions import WrongSecret
+from .protocols.telegram import TelegramProtocol
 
 
 class Dispatcher:
-    def __init__(self, loop, real_protocol):
+    def __init__(self, loop, real_protocol, secret):
         self.loop = loop
+        self.secret = secret
         self.real_protocol = real_protocol
         self.telegram_protocol = None
         self.telegram_transport = None
@@ -15,14 +16,20 @@ def __init__(self, loop, real_protocol):
 
     async def handle(self, raw_data: bytes):
         if not self.encryptor or not self.decryptor:
-            self.decryptor = Decryptor(raw_data, secret)
-            self.encryptor = Encryptor(raw_data, secret)
+            self.decryptor = Decryptor(raw_data, self.secret)
+            self.encryptor = Encryptor(raw_data, self.secret)
             self.decoded_data = await self.decryptor(raw_data)
             server_id = int.from_bytes(self.decoded_data[60:62], 'little') - 1
             if server_id < 0 or server_id > len(telegram_ips):
-                raise WrongServerID(f'Got id {server_id}')
+                raise WrongSecret(f'Got id {server_id}')
             self.telegram_server_ip = telegram_ips[server_id]
-            assert self.decoded_data[56:60] == mtproto_footprint or self.decoded_data[56:60] == obf_footprint, self.decoded_data[56:60]
+            if self.decoded_data[56:60] == mtproto_footprint:
+                pass
+            elif self.decoded_data[56:60] == obf_footprint:
+                # TODO dd prefix (padding)
+                pass
+            else:
+                raise Exception('Unknown Protocol')
         if not self.telegram_protocol or not self.telegram_transport:
             self.telegram_transport, self.telegram_protocol = await self.loop.create_connection(lambda: TelegramProtocol(self.loop, self),
                 self.telegram_server_ip, 443)

amtproxy/encryption/crypto.py

@@ -1,6 +1,6 @@
 from hashlib import sha256
 
-from amtproxy.encryption.aes import AES_CTR
+from .aes import AES_CTR
 
 
 class DecryptionKey:

amtproxy/exceptions.py

@@ -4,5 +4,5 @@ class ProxyBaseException(Exception):
     pass
 
 
-class WrongServerID(ProxyBaseException):
+class WrongSecret(ProxyBaseException):
     pass

amtproxy/protocols/__init__.py

@@ -1,2 +0,0 @@
-from .real import RealProtocol
-from .telegram import TelegramProtocol

amtproxy/protocols/real.py

@@ -1,29 +1,35 @@
 import asyncio
-from amtproxy.dispatcher import Dispatcher
+from ..dispatcher import Dispatcher
 
 
 class RealProtocol(asyncio.Protocol):
-    def __init__(self, loop, *args, **kwargs):
+    transport: asyncio.transports.BaseTransport
+    task: asyncio.Task = None
+
+    def __init__(self, loop, secret, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.loop = loop
-        self.dispatcher = Dispatcher(loop, self)
+        self.dispatcher = Dispatcher(loop, self, secret)
 
     def connection_made(self, transport):
         self.transport = transport
 
+
     def data_received(self, data):
-        task = asyncio.Task(self.async_handler(data), loop=self.loop)
+        self.task = asyncio.Task(self.async_handler(data), loop=self.loop)
 
 
     @asyncio.coroutine
     def async_handler(self, data):
         return (yield from self.dispatcher.handle(data))
 
     def connection_lost(self, exc) -> None:
+        if self.task:
+            self.task.cancel()
         try:
             self.dispatcher.telegram_transport.close()
         except AttributeError:
             pass
 
     def close_connection(self):
-        self.transport.close()
+        self.transport.close()

amtproxy/protocols/telegram.py

@@ -1,8 +1,11 @@
 import asyncio
+from asyncio.transports import BaseTransport
 
 
 class TelegramProtocol(asyncio.Protocol):
-    
+    transport: BaseTransport
+    task: asyncio.Task = None
+
     def __init__(self, loop, dispatcher, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.loop = loop
@@ -12,11 +15,13 @@ def connection_made(self, transport):
         self.transport = transport
 
     def data_received(self, data):
-        task = asyncio.Task(self.async_handler(data), loop=self.loop)
+        self.task = asyncio.Task(self.async_handler(data), loop=self.loop)
 
     @asyncio.coroutine
     def async_handler(self, data):
         return (yield from self.dispatcher.reverse_handle(data))
 
     def connection_lost(self, exc) -> None:
-        self.dispatcher.real_protocol.close_connection()
+        if self.task:
+            self.task.cancel()
+        self.dispatcher.real_protocol.close_connection()

amtproxy/server.py

@@ -1,22 +1,23 @@
 import asyncio
-from amtproxy.protocols import RealProtocol
 
+from .protocols.real import RealProtocol
 
-async def main(port=8080):
-    loop = asyncio.get_running_loop()
 
+async def main(port=8080, secret='1234567890abc1234567890abcef1234'):
+    loop = asyncio.get_running_loop()
+    print(f'MTProxy is starting on port {port}\nWith secret {secret}')
     server = await loop.create_server(
-        lambda: RealProtocol(loop),
+        lambda: RealProtocol(loop, secret),
         '0.0.0.0', port)
 
     async with server:
         await server.serve_forever()
 
 
-def start(port=None):
-    asyncio.run(main(port))
+def start_server(port=None, secret=None):
+    args = {'port': port, 'secret': secret}
+    asyncio.run(main(**{k: v for k, v in args.items() if v}))
 
 
 if __name__ == '__main__':
     asyncio.run(main())
-

setup.py

@@ -11,7 +11,7 @@
 
 setuptools.setup(
     name="aMTProxy",
-    version="0.1.5",
+    version="0.2",
     author="Yury (Yurzs)",
     author_email="dev@best-service.online",
     description="Async python MTProxy server",