ServerJar/client.py

import argparse
import asyncio
import queue
import ssl
import sys
import threading
import socket
import time
import traceback
from prompt_toolkit import Application
from prompt_toolkit.layout import Layout, HSplit
from prompt_toolkit.widgets import TextArea
from prompt_toolkit.key_binding import KeyBindings
from prompt_toolkit.styles import Style
from prompt_toolkit.filters import has_focus
from prompt_toolkit.shortcuts import clear as ptk_clear

version = "Beta-1"


class ServerJarClient(Application):
    def __init__(self, **kwargs):
        super().__init__(**kwargs, mouse_support=True)

        # Text style
        self.style = Style.from_dict({
            "log": "bg:#000000 #ffffff",
            "input": "bg:#222222 #ffffff",
            "separator-area": "bg:#000000 #ffffff",
            "message-area": "bg:#111111 #ffffff"
        })
        # Socket
        self.sock = None

        # Args
        self.args = None

        # event
        self.kb = KeyBindings()

        # Areas
        # self.log_lines = []

        # self.log_control = FormattedTextControl(
        #     text=lambda: ANSI("".join(self.log_lines))
        # )

        self.log_area = TextArea(
            style="class:log",
            wrap_lines=True,
        )

        self.separator_area = TextArea(text="=" * 10 + " Enter Command Here " + "=" * 10, height=1,
                                       style="class:separator-area")
        self.message_area = TextArea(height=1, multiline=False, style="class:message-area")
        self.input_area = TextArea(height=1, prompt="> ", style="class:input", multiline=False)

        self.layout = Layout(HSplit([
            self.log_area,
            self.separator_area,
            self.message_area,
            self.input_area,
        ]))

        # Thread
        self.sock_lock = threading.Lock()
        self.closing_event = threading.Event()
        self.client_thread = threading.Thread(target=self.client, daemon=True)
        self.connect_event = threading.Event()
        self.disconnect_event = threading.Event()

        # Queue
        self.incoming = queue.Queue()

        # Command History
        self.cmds = []
        self.current_index = None
        self.start_history_flag = False

        @self.kb.add("c-c")
        def closing_kb(event):
            self.shutdown("Ctrl-C (Stopped by user)")

        @self.kb.add("up", filter=has_focus(self.input_area))
        def get_old_cmd(event):
            # check if there's no old command available in the command history list
            if len(self.cmds) < 2:
                return

            # Save entered command to history list
            if not self.start_history_flag:
                self.insert_new_cmd_to_history(self.get_input_area_output())
                self.start_history_flag = True

            # Set current_index to 1 if it's not initiated yet
            if self.current_index is None:
                self.current_index = 0

            # Avoid IndexError when it's the last one command
            if self.current_index >= len(self.cmds)-1:
                return

            self.current_index += 1

            old_cmd = self.cmds[self.current_index]

            self.set_input_area_text(old_cmd)

        @self.kb.add("down", filter=has_focus(self.input_area))
        def get_new_cmd(event):
            # Check if there's no old command available in the command history list
            if len(self.cmds) < 2:
                return

            # Get new command only working when current_index greater then 0
            if self.current_index is None:
                return

            # Avoid IndexError when it's the last one command
            if self.current_index-1 < 0:
                return

            self.current_index -= 1

            new_cmd = self.cmds[self.current_index]

            self.set_input_area_text(new_cmd)

        @self.kb.add("enter", filter=has_focus(self.input_area))
        def enter_kb(event):
            # Disable history flag
            self.start_history_flag = False

            cmd = self.input_area.text
            self.input_area.text = ""

            # Save command
            self.insert_new_cmd_to_history(cmd)

            if cmd == "_exit":
                self.shutdown("_exit command detected")
                return

            # if re.match(r"^_[A-Za-z0-9]+(?:$|_.*)", cmd):
            exit_flag = self.command_parser(cmd)

            if exit_flag:
                return

            with self.sock_lock:
                s = self.sock

            if s:
                try:
                    s.sendall((cmd + "\n").encode("utf-8"))
                except OSError as e:
                    self._err(f"Send failed: {e}\n")
            else:
                self._err("The remote server is not connected yet.")

        @self.kb.add("c-w", filter=has_focus(self.input_area))
        def focus_log_area(event):
            self.layout.focus(self.log_area)
            self.display_message("Now focus at log area.")

        @self.kb.add("c-w", filter=has_focus(self.log_area))
        def focus_log_area(event):
            self.layout.focus(self.input_area)
            self.display_message("Now focus at input area.")

        self.key_bindings = self.kb
        self.full_screen = True

        # Host and Port
        self.host = None
        self.port = None

    def command_parser(self, command):
        def connect_to_server(host, port):
            # update target
            self.host = host
            self.port = port

            # trigger connection
            self.disconnect_event.clear()
            self.connect_event.set()

            return True

        def disconnect_from_server(cmd):
            self._log("Disconnecting...")
            self.disconnect_event.set()
            self.connect_event.clear()

            with self.sock_lock:
                s = self.sock

            if s:
                try:
                    s.shutdown(socket.SHUT_RDWR)
                except Exception as e:
                    self._err("An error occurred while shutting down sock: " + str(e))
                try:
                    s.close()
                except Exception as e:
                    self._err("An error occurred while closing the socket: {}".format(e))
            else:
                self._err("The remote server is not connected yet.")

            return True

        def connect_to_server_parser(cmd):
            target = cmd[3:].strip()
            try:
                host, port_str = target.split(":", 1)
                host = host.strip()
                port = int(port_str.strip())

                if not host:
                    raise ValueError("empty host")
            except Exception as _:
                self._err("Usage: _c host:port")
                return True

            connect_to_server(host, port)

            return True

        def _shutdown(cmd):
            self.shutdown("_exit command detected")
            return True

        def _top(cmd):
            self.log_area.buffer.cursor_position = 0
            self.invalidate()
            return True

        def _bottom(cmd):
            self.log_area.buffer.cursor_position = len(self.log_area.buffer.text)
            self.invalidate()
            return True

        def _version(cmd):
            self._log("ServerJar Client Version {}".format(version))
            return True

        cmd_map = {
            "_exit": _shutdown,
            "_c": connect_to_server_parser,
            "_d": disconnect_from_server,
            "_top": _top,
            "_bottom": _bottom,
            "_version": _version,
        }

        for cmd in cmd_map.keys():
            if command.startswith(cmd):
                return_flag = cmd_map[cmd](command)
                return return_flag

        # self._err("Unknown command '%s'" % command)

        return False

    def display_message(self, message):
        self.message_area.text = message

    def get_input_area_output(self):
        return self.input_area.text

    def set_input_area_text(self, text):
        self.input_area.text = text

    def insert_new_cmd_to_history(self, cmd):
        self.cmds.insert(0, cmd)

    class ServerInfoInvalidException(Exception):
        def __init__(self, message, **kwargs):
            super().__init__()
            self.msg = message

        def __str__(self):
            return self.msg

    def arguments_parser(self):
        parser = argparse.ArgumentParser()

        parser.add_argument("-p", "--port", type=int, help="Port number", required=True)
        parser.add_argument('-host', '--host', type=str, help="Hostname", required=True)
        parser.add_argument('-no-tls', '--no-tls', type="store_true", help="Enable TLS support")

        args = parser.parse_args()

        return args

    def shutdown(self, reason=""):
        if self.closing_event.is_set():
            return

        self.closing_event.set()

        self._log(f"Shutting down for reason: {reason}")

        with self.sock_lock:
            s = self.sock
            if s:
                try:
                    s.close()
                except Exception as e:
                    self._err(f"Unable to close socket: {e}")
                    pass

        # Exit ui event loop
        self.full_exit()

    # @staticmethod
    # def clear_screen():
    #     os.system("cls" if os.name == "nt" else "clear")

    def log(self, message):
        self.incoming.put(f"{message}")

    def _log(self, message):
        # Nothing change
        self.incoming.put(f"[client] {message}")
        # self.display_message(f"{message}")

    def _err(self, message):
        # WIP... (display text as red color if the log is an error message)
        self.incoming.put(f"[client|err] {message}")
        # self.display_message(f"ERROR: {message}")

    def _warn(self, message):
        # WIP... (Display text as yellow color if the log is a warning message)
        self.incoming.put(f"[client|warn] {message}")
        # self.display_message(f"WARNING: {message}")

    def full_exit(self):
        self.exit()
        sys.exit()

    async def consume_incoming(self):
        loop = asyncio.get_running_loop()
        while True:
            msg = await loop.run_in_executor(None, self.incoming.get)

            # self.log_lines.append(msg)
            #
            # if len(self.log_lines) > 2000:
            #     self.log_lines = self.log_lines[-2000:]

            if len(self.log_area.text) > 0:
                self.log_area.text += "\n" + msg
            else:
                self.log_area.text += msg

            if len(self.log_area.text) > 300_000:
                self.log_area.text = "New log start here.\n" + self.log_area.text[-250_000:]

            self.log_area.buffer.cursor_position = len(self.log_area.buffer.text)

            self.invalidate()

    def client(self):
        ptk_clear()


        while not self.closing_event.is_set():
            self._log("Type _c host:port to connect. (_d to disconnect)")
            self.connect_event.wait()

            if self.closing_event.is_set():
                break

            if not self.host or not self.port:
                self._err("No host/port set. Usage: _c host:port")
                self.connect_event.clear()
                continue

            try:
                self._log(f"Connecting to {self.host}:{self.port} ...")

                # Create connect
                if self.args.no_tls:
                    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                    s.connect((self.host, self.port))
                else:
                    raw = socket.create_connection((self.host, self.port))
                    context = ssl.create_default_context()
                    s = context.wrap_socket(raw, server_hostname=self.host)

                s.connect((self.host, self.port))

                with self.sock_lock:
                    self.sock = s

                self._log("Remote socket server connected [HOST: {}, PORT: {}]".format(self.host, self.port))

                buffer = ""
                while True:
                    # Receive remote server broadcast message and display it on log area
                    data = s.recv(4096)
                    if not data:
                        raise ConnectionError("Server closed")
                    buffer += data.decode("utf-8", errors="replace")
                    while "\n" in buffer:
                        line, buffer = buffer.split("\n", 1)
                        # ### Use normal log method ###
                        self.log(line)

            except (ConnectionError, OSError) as e:
                if not self.closing_event.is_set():
                    self._warn(f"Disconnected: {e}, retrying...")
                    time.sleep(1)
            except KeyboardInterrupt:
                self._log("Exiting...")
                break
            except Exception as e:
                self._err(f"Unhandled exception: {e}")
                self._err(f"{traceback.format_exc()}")
            finally:
                with self.sock_lock:
                    try:
                        if self.sock:
                            self._log("Closing remote connection (From {}:{})...".format(self.host, self.port))
                            self.sock.close()
                    except Exception as e:
                        self._err(f"Unable to close socket: {e}")
                        pass
                    self.sock = None

                # reset flags
                self.disconnect_event.clear()
                self.connect_event.clear()


    def startup(self):
        self.args = self.arguments_parser()
        self.layout.focus(self.input_area)
        asyncio.create_task(self.consume_incoming())

        self.client_thread.start()


if __name__ == "__main__":
    app = ServerJarClient()
    app.run(pre_run=app.startup)