Merge pull request #14 from Code-Partners/feat-pipe-protocol

Feat pipe protocol

Author: SergeyCP

protocols/pipe_protocol/__init__.py

@@ -0,0 +1,79 @@
+import io
+import platform
+import typing
+
+from common.exceptions import SmartInspectError
+from connections.builders import ConnectionsBuilder
+from formatters import BinaryFormatter
+from packets.packet import Packet
+from protocols.protocol import Protocol
+
+
+class PipeProtocol(Protocol):
+    _BUFFER_SIZE: int = 0x2000
+    _CLIENT_BANNER_TEMPLATE: str = "SmartInspect Python Library v{}\n"
+    _PIPE_NAME_PREFIX: str = "\\\\.\\pipe\\"
+
+    def __init__(self):
+        super().__init__()
+        self._formatter: BinaryFormatter = BinaryFormatter()
+        self._stream: typing.Optional[io.RawIOBase, io.BufferedIOBase] = None
+        self._pipe_name: str = ""
+        self._load_options()
+
+    @staticmethod
+    def _get_name() -> str:
+        return "pipe"
+
+    def _do_handshake(self) -> None:
+        self._read_server_banner()
+        self._send_client_banner()
+
+    def _read_server_banner(self) -> None:
+        answer = self._stream.readline().strip()
+        if not answer:
+            raise SmartInspectError("Could not read server banner correctly: " +
+                                    "Connection has been closed unexpectedly")
+
+    def _send_client_banner(self) -> None:
+        from smartinspect import SmartInspect
+        si_version = SmartInspect.get_version()
+        self._stream.write(self._CLIENT_BANNER_TEMPLATE.format(si_version).encode("UTF-8"))
+        self._stream.flush()
+
+    def _is_valid_option(self, option_name: str) -> bool:
+        return (option_name == "pipename"
+                or super()._is_valid_option(option_name))
+
+    def _build_options(self, builder: ConnectionsBuilder) -> None:
+        super()._build_options(builder)
+        builder.add_option("pipename", self._pipe_name)
+
+    def _load_options(self) -> None:
+        super()._load_options()
+        self._pipe_name = self._get_string_option("pipename", "smartinspect")
+
+    def _internal_connect(self) -> None:
+        if platform.system() != "Windows":
+            raise SmartInspectError("Pipe Protocol is only supported on Windows")
+
+        filename = self._PIPE_NAME_PREFIX + self._pipe_name
+        try:
+            self._stream = open(filename, "w+b", buffering=self._BUFFER_SIZE)
+        except Exception as e:
+            raise SmartInspectError(f"\nThere was a connection error. \n"
+                                    f"Check if pipe with name <{filename}> exists\n"
+                                    f"Your system returned: {type(e).__name__} - {str(e)}")
+        self._do_handshake()
+        self._internal_write_log_header()
+
+    def _internal_write_packet(self, packet: Packet) -> None:
+        self._formatter.format(packet, self._stream)
+        self._stream.flush()
+
+    def _internal_disconnect(self) -> None:
+        if self._stream:
+            try:
+                self._stream.close()
+            finally:
+                self._stream = None

protocols/pipe_protocol/pipe_protocol.py

@@ -4,6 +4,7 @@
 from protocols.cloud.cloud_protocol import CloudProtocol
 from protocols.file_protocol.file_protocol import FileProtocol
 from protocols.memory_protocol import MemoryProtocol
+from protocols.pipe_protocol.pipe_protocol import PipeProtocol
 from protocols.protocol import Protocol
 from common.exceptions import SmartInspectError
 from protocols.tcp_protocol import TcpProtocol
@@ -18,6 +19,7 @@ class ProtocolFactory:
         "file": FileProtocol,
         "text": TextProtocol,
         "mem": MemoryProtocol,
+        "pipe": PipeProtocol,
     }
     __PROTOCOL_NOT_FOUND = "The requested protocol is unknown"
     __lock = threading.Lock()

protocols/protocol_factory.py

@@ -12,13 +12,12 @@
 
 
 class TcpProtocol(Protocol):
-    __BUFFER_SIZE = 0x2000
-    __CLIENT_BANNER = bytearray(f"SmartInspect Python Library v\n", encoding="UTF-8")
-    # __CLIENT_BANNER = bytearray(f"SmartInspect Python Library v{SmartInspect.get_version()}\n", encoding="UTF-8")
-    __ANSWER_BUFFER_SIZE = 0x2000
-    _hostname = "127.0.0.1"
-    _timeout = 30000
-    _port = 4228
+    __BUFFER_SIZE: int = 0x2000
+    __CLIENT_BANNER_TEMPLATE: str = "SmartInspect Python Library v{}\n"
+    __ANSWER_BUFFER_SIZE: int = 0x2000
+    _hostname: str = "127.0.0.1"
+    _timeout: int = 30000
+    _port: int = 4228
 
     def __init__(self):
         super().__init__()
@@ -59,7 +58,9 @@ def _read_server_banner(self) -> None:
                                     "Connection has been closed unexpectedly")
 
     def _send_client_banner(self) -> None:
-        self.__stream.write(self.__CLIENT_BANNER)
+        from smartinspect import SmartInspect
+        si_version = SmartInspect.get_version()
+        self.__stream.write(self.__CLIENT_BANNER_TEMPLATE.format(si_version).encode("UTF-8"))
         self.__stream.flush()
 
     def _internal_connect(self):

protocols/tcp_protocol.py

@@ -68,9 +68,9 @@ def set_resolution(self, resolution: ClockResolution) -> None:
         if isinstance(resolution, ClockResolution):
             self.__resolution = resolution
 
-    @property
-    def version(self) -> str:
-        return self.__VERSION
+    @classmethod
+    def get_version(cls) -> str:
+        return cls.__VERSION
 
     @property
     def hostname(self) -> str: