cubething_psync 0.2.1.dev1__tar.gz → 0.3.0__tar.gz

{cubething_psync-0.2.1.dev1/src/cubething_psync.egg-info → cubething_psync-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cubething_psync
-Version: 0.2.1.dev1
+Version: 0.3.0
 Summary: Simple project synchronization tool.
 Author-email: ada mandala <ada@cubething.dev>
 License-Expression: MIT

{cubething_psync-0.2.1.dev1 → cubething_psync-0.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "cubething_psync"
-version = "0.2.1.dev1"
+version = "0.3.0"
 description = "Simple project synchronization tool."
 authors = [{name = "ada mandala", email="ada@cubething.dev"}]
 readme = "README.md"
@@ -26,6 +26,18 @@ build-backend = "setuptools.build_meta"
 [tool.setuptools.packages.find]
 where= ["src"]
 
+[dependency-groups]
+dev = [
+    "basedpyright>=1.35.0",
+    "pytest>=9.0.1",
+    "ruff>=0.14.7",
+    "testcontainers>=4.13.3",
+]
+
 [project.scripts]
 psync-server = "server.main:main"
 psync-client = "client.main:main"
+
+[pytest]
+testpaths="test"
+addopts = ["--import-mode=importlib"]

cubething_psync-0.3.0/src/client/args.py

@@ -0,0 +1,190 @@
+import argparse
+from dataclasses import dataclass, field
+import hashlib
+import logging
+import os
+from os.path import basename
+from pathlib import Path
+import shlex
+from common.data import deserialize_env
+
+
+@dataclass
+class Args:
+    """
+    Client arguments.
+    """
+
+    target_path: str
+    """
+    ``--path -p <target_path>``
+
+    **Required.** Path to the target executable.
+    """
+
+    assets: list[str] = field(default_factory=list)
+    """
+    ``--assets -A <env>``
+
+    Extra files or directories to be synced to the destination path.
+    """
+
+    env: dict[str, str] = field(default_factory=dict)
+    """
+    ``--env -e <env>``
+
+    Space separated list of environment variables which will be passed to the executable.
+    """
+
+    args: list[str] = field(default_factory=list)
+    """
+    ``--args -a <args>``
+
+    Arguments passed to the executable.
+    """
+
+    server_ip: str = os.environ.get("PSYNC_SERVER_IP", "127.0.0.1")
+    """
+    environ: ``PSYNC_SERVER_IP``
+
+    Server IP address.
+    """
+
+    server_port: int = int(os.environ.get("PSYNC_SERVER_PORT", "5000"))
+    """
+    environ: ``PSYNC_SERVER_PORT``
+
+    Server port.
+    """
+
+    server_ssh_port: int = int(os.environ.get("PSYNC_SSH_PORT", "5022"))
+    """
+    environ: ``PSYNC_SSH_PORT``
+
+    SSH port on the server host. Client must be authenticated with a shared public key.
+    """
+
+    ssh_args: str = os.environ.get("PSYNC_SSH_ARGS", "-l psync")
+    """
+    environ: ``PSYNC_SSH_ARGS``
+
+    Arguments passed to SSH. Under the hood, psync runs
+    ``rsync -e "/usr/bin/ssh {PSYNC_SSH_ARGS} -p {PSYNC_SSH_PORT}"``
+    """
+
+    server_dest: str = os.environ.get("PSYNC_SERVER_DEST", "/home/psync")
+    """
+    environ: ``PSYNC_SERVER_DEST``
+
+    Base path on the server where the files should be synced.
+    """
+
+    ssl_cert_path: str = os.environ.get(
+        "PSYNC_CERT_PATH", "~/.local/share/psync/cert.pem"
+    )
+    """
+    environ: ``PSYNC_CERT_PATH``
+
+    Public SSL certificate used to trust the psync server.
+    """
+
+    client_origin: str = os.environ.get("PSYNC_CLIENT_ORIGIN", "127.0.0.1")
+    """
+    environ: ``PSYNC_CLIENT_ORIGIN``
+
+    Domain name. Should match the origins set in the server's ``PSYNC_ORIGINS``
+    variable.
+    """
+
+    def project_hash(self) -> str:
+        """
+        Hash value generated from the target path. Used as the directory name for the project.
+        """
+        return hashlib.blake2s(self.target_path.encode(), digest_size=8).hexdigest()
+
+    def rsync_url(self) -> str:
+        """
+        {server_ip}:{server_dest}/{project_hash}
+        """
+        return f"{self.server_ip}:{self.server_dest}/{self.project_hash()}/"
+
+    def destination_path(self) -> Path:
+        """
+        {server_dest}/{project_hash}/{basename(target_path)}
+        """
+        return Path(self.server_dest) / self.project_hash() / basename(self.target_path)
+
+
+parser = argparse.ArgumentParser(
+    prog="psync-client",
+    usage="""\
+Client for the psync server.
+
+In addition to the options below, the client is configurable through environment
+variables.
+
+Variable            | Default
+--------------------+-------------------------------
+PSYNC_SERVER_IP     | 127.0.0.1
+PSYNC_SERVER_PORT   | 5000
+PSYNC_SSH_PORT      | 5022
+PSYNC_SERVER_DEST   | /home/psync/
+PSYNC_SSH_ARGS      | -l psync
+PSYNC_CERT_PATH     | ~/.local/share/psync/cert.pem
+PSYNC_CLIENT_ORIGIN | 127.0.0.1
+
+SSH arguments will be append with "-p {PSYNC_SSH_PORT}"
+""",
+)
+_action = parser.add_argument(
+    "--path",
+    "-p",
+    required=True,
+    help="Path to the target exectuable.",
+)
+_action = parser.add_argument(
+    "--assets",
+    "-A",
+    nargs="+",
+    help="Extra files or directories to be synced to the destination path.",
+)
+_action = parser.add_argument(
+    "--env",
+    "-e",
+    help="Environment variables to set in the remote execution environment. Variables must be space-sepated or double-quoted.",
+)
+_action = parser.add_argument(
+    "--args", "-a", help="Arguments passed to the executable."
+)
+
+
+def parse_args() -> Args:
+    args = vars(parser.parse_args())
+
+    target_path = str(args.get("path"))
+    target_path = Path(target_path)
+    if not target_path.is_file():
+        logging.error(f"Could not file at {target_path}")
+        exit(1)
+
+    extra: list[str] = []
+    extra_raw = args.get("extra")
+    if extra_raw is not None:
+        extra = extra_raw  # pyright: ignore[reportAny]
+
+    client_args: list[str] = []
+    raw_args = args.get("args")
+    if raw_args is not None:
+        client_args = shlex.split(str(raw_args))  # pyright: ignore[reportAny]
+
+    env: dict[str, str] = dict()
+    raw_env = args.get("env")
+    if raw_env is not None:
+        env = deserialize_env(f"env='{raw_env}'")
+
+    return Args(
+        target_path=str(target_path),
+        assets=extra or [],
+        env=env,
+        args=client_args,
+    )

{cubething_psync-0.2.1.dev1 → cubething_psync-0.3.0}/src/client/main.py

@@ -3,13 +3,13 @@ psync client
 """
 
 import asyncio
-import hashlib
 import os
-import pathlib
+from pathlib import Path
 import signal
 import ssl
 import subprocess
 import websockets
+from websockets.typing import Origin
 from common.data import (
     ErrorResp,
     ExitResp,
@@ -17,18 +17,13 @@ from common.data import (
     LogResp,
     OkayResp,
     OpenReq,
+    SetPidResp,
     deserialize,
     serialize,
 )
 import logging
 from common.log import InterceptHandler
 from client.args import (
-    SERVER_IP,
-    SERVER_PORT,
-    SERVER_SSH_PORT,
-    USER,
-    SSL_CERT_PATH,
-    SERVER_DEST,
     Args,
     parse_args,
 )
@@ -38,69 +33,54 @@ class PsyncClient:
     """
     The primary interface for psync. The client CLI allows users to sync files with
     rsync, then execute them remotely while receiving the logs.
-
-    CLI arguments: ::
-          -h, --help            show this help message and exit
-          --path, -p PATH       Path to the target exectuable.
-          --extra, -E EXTRA [EXTRA ...]
-                                Extra files or directories to be synced to the destination path.
-          --env, -e ENV         Environment variables to set in the remote execution environment. Variables
-                                must be space-sepated or double-quoted.
-          --args, -a ARGS       Arguments with which to run the remote executable.
-
-    Environment configuration:
-        PSYNC_SERVER_IP: The IP address of the server instance.
-            Default: 127.0.0.1
-        PSYNC_SERVER_PORT: The port of the server instance.
-            Default: 5000
-        PSYNC_SSH_PORT: The server instance's SSH port.
-            Default: 5022
-        PSYNC_SSH_USER: The server instance's SSH user.
-            Default: psync
-        PSYNC_CERT_PATH: Path to the SSL certificate. Used to trust self-signed certs. Should
-            match the server's certificate.
-            Default: ~/.local/share/psync/cert.pem
     """
 
-    __args: list[str]
-    __env: dict[str, str]
-    __path: pathlib.Path
-    __quit: bool = False
+    args: Args
+    pid: int | None = None
+    """ID of the current connection."""
     __force_exit: bool = False
 
-    def __init__(self, args: list[str], env: dict[str, str], path: pathlib.Path):
-        self.__args = args
-        self.__env = env
-        self.__path = path
+    def __init__(self, args: Args):
+        self.args = args
 
     def __mk_handler(self, ws: websockets.ClientConnection):
         async def inner():
             if not self.__force_exit:
                 logging.info("Gracefully shutting down...")
                 self.__force_exit = True
-                await ws.send(serialize(KillReq()))
+                if self.pid is not None:
+                    await ws.send(serialize(KillReq(pid=self.pid)))
                 await ws.close()
                 asyncio.get_event_loop().stop()
+                raise SystemExit(130)
             else:
                 logging.warning("Got second SIGINT, shutting down")
                 asyncio.get_event_loop().stop()
-                raise Exception("Forced shutdown")
+                raise SystemExit(1)
 
         return lambda: asyncio.create_task(inner())
 
     async def run(self):
         """Run the client instance."""
         ssl_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
-        ssl_ctx.load_verify_locations(pathlib.Path(SSL_CERT_PATH).expanduser())
+        ssl_ctx.load_verify_locations(Path(self.args.ssl_cert_path).expanduser())
         ssl_ctx.check_hostname = False  # not ideal
         async with websockets.connect(
-            f"wss://{SERVER_IP}:{SERVER_PORT}", ssl=ssl_ctx
+            f"wss://{self.args.server_ip}:{self.args.server_port}",
+            ssl=ssl_ctx,
+            origin=Origin(f"wss://{self.args.client_origin}"),
         ) as ws:
             asyncio.get_event_loop().add_signal_handler(
                 signal.SIGINT, self.__mk_handler(ws)
             )
             await ws.send(
-                serialize(OpenReq(path=self.__path, env=self.__env, args=self.__args))
+                serialize(
+                    OpenReq(
+                        path=self.args.destination_path(),
+                        env=self.args.env,
+                        args=self.args.args,
+                    )
+                )
             )
             async for data in ws:
                 if isinstance(data, bytes):
@@ -115,7 +95,7 @@ class PsyncClient:
                         f"Failed to deserialize message '{msg}' with error '{e}'"
                     )
                     await ws.close()
-                    exit(1)
+                    raise Exception(e)
 
                 match resp:
                     case LogResp():
@@ -123,39 +103,42 @@ class PsyncClient:
                     case ErrorResp():
                         logging.error(f"Received server error: {resp.msg}")
                         await ws.close()
-                        exit(1)
+                        raise Exception(resp.msg)
                     case ExitResp():
                         logging.info(f"Exiting with code {resp.exit_code}")
                         await ws.close()
-                        exit(resp.exit_code)
+                        raise SystemExit(resp.exit_code)
+                    case SetPidResp():
+                        logging.info(f"Remote PID = {resp.pid}")
+                        self.pid = resp.pid
                     case OkayResp():
-                        logging.info("Running exectuable...")
+                        logging.info("OK.")
                     case _:
                         logging.warning(f"Got unknown request {resp}")
 
 
-def __rsync(project_hash: str, args: Args):
+def __rsync(args: Args):
     """Runs rsync."""
-    url = f"{SERVER_IP}:{SERVER_DEST}/{project_hash}/"
     rsync_args = [
         "rsync",
         "-avzr",
         "-e",
-        f"/usr/bin/ssh -l {USER} -p {str(SERVER_SSH_PORT)}",
+        f"/usr/bin/ssh {args.ssh_args} -p {str(args.server_ssh_port)}",
         "--progress",
         "--mkpath",
         args.target_path,
-        *args.extra,
-        url,
+        *args.assets,
+        args.rsync_url(),
     ]
     logging.info(" ".join(rsync_args))
     p = subprocess.run(rsync_args)
     if p.returncode != 0:
-        logging.error(f"Rsync failed with exit code {p.returncode}")
-        exit(1)
+        msg = f"Rsync failed with exit code {p.returncode}"
+        logging.error(msg)
+        raise Exception(msg)
 
 
-def main():
+def main(args: Args | None = None):
     """
     The main executable.
     Sync project files with rsync, then run the client.
@@ -163,15 +146,15 @@ def main():
     log_level = os.environ.get("PSYNC_LOG", "INFO").upper()
     logging.basicConfig(handlers=[InterceptHandler()], level=log_level, force=True)
 
-    project_hash = hashlib.blake2s(os.getcwd().encode(), digest_size=8).hexdigest()
-    args = parse_args()
-    __rsync(project_hash, args)
+    args = parse_args() if args is None else args
+    __rsync(args)
 
-    dest_path = pathlib.Path(
-        f"{SERVER_DEST}/{project_hash}/{os.path.basename(args.target_path)}"
-    )
-    client = PsyncClient(args=args.args, env=args.env, path=dest_path)
-    asyncio.run(client.run())
+    try:
+        asyncio.run(PsyncClient(args).run())
+    except SystemExit as e:
+        exit(e.code)
+    except Exception:
+        exit(1)
 
 
 if __name__ == "__main__":

{cubething_psync-0.2.1.dev1 → cubething_psync-0.3.0}/src/common/data.py

@@ -15,6 +15,7 @@ class Mode(Enum):
 class ReqKind(Enum):
     Open = "open"
     Kill = "kill"
+    HealthCheck = "hc"
 
 
 class RespKind(Enum):
@@ -22,6 +23,7 @@ class RespKind(Enum):
     Error = "error"
     Exit = "exit"
     Okay = "ok"
+    SetPid = "set_pid"
 
 
 @dataclass
@@ -34,6 +36,7 @@ class OpenReq:
 
 @dataclass
 class KillReq:
+    pid: int
     kind: ReqKind = ReqKind.Kill
 
 
@@ -60,28 +63,44 @@ class OkayResp:
     kind: RespKind = RespKind.Okay
 
 
-Req = OpenReq | KillReq
-Resp = LogResp | ExitResp | ErrorResp | OkayResp
+@dataclass
+class HealthCheckReq:
+    kind: ReqKind = ReqKind.HealthCheck
+
+
+@dataclass
+class SetPidResp:
+    pid: int
+    kind: RespKind = RespKind.SetPid
+
+
+Req = OpenReq | KillReq | HealthCheckReq
+Resp = LogResp | ExitResp | ErrorResp | OkayResp | SetPidResp
 
 
 def serialize(msg: Req | Resp) -> str:
+    value = msg.kind.value
     match msg:
         case OpenReq():
             args = " ".join(msg.args)
             env: list[str] = []
             for key, value in msg.env.items():
                 env.append(f'{key}="{value}"')
-            return f"open path='{msg.path}' args='{args}' env='{' '.join(env)}'"
+            return f"{value} path='{msg.path}' args='{args}' env='{' '.join(env)}'"
         case KillReq():
-            return "kill"
+            return f"{value} {msg.pid}"
+        case SetPidResp():
+            return f"{value} {msg.pid}"
         case LogResp():
-            return f"log {msg.msg}"
+            return f"{value} {msg.msg}"
         case ExitResp():
-            return f"exit {msg.exit_code}"
+            return f"{value} {msg.exit_code}"
         case ErrorResp():
-            return f"error {msg.msg}"
+            return f"{value} {msg.msg}"
         case OkayResp():
-            return "okay"
+            return f"{value}"
+        case HealthCheckReq():
+            return f"{value}"
 
 
 path_expr = re.compile(r"path='([^']+)'")
@@ -106,13 +125,8 @@ def deserialize(msg: str) -> Req | Resp:
     logging.debug(f"Got message {msg}")
     try:
         [kind, rest] = msg.split(" ", 1)
-    except Exception as e:
-        if msg == "okay":
-            return OkayResp()
-        elif msg == "kill":
-            return KillReq()
-        else:
-            raise e
+    except Exception:
+        kind = msg.strip()
 
     match kind:
         case ReqKind.Open.value:
@@ -134,12 +148,20 @@ def deserialize(msg: str) -> Req | Resp:
 
             return OpenReq(path=Path(path), args=args, env=env)
 
+        case ReqKind.Kill.value:
+            return KillReq(int(rest))
         case RespKind.Log.value:
-            return LogResp(msg=rest)
+            return LogResp(rest)
         case RespKind.Exit.value:
-            return ExitResp(exit_code=rest)
+            return ExitResp(rest)
         case RespKind.Error.value:
-            return ErrorResp(msg=rest)
+            return ErrorResp(rest)
+        case RespKind.Okay.value:
+            return OkayResp()
+        case ReqKind.HealthCheck.value:
+            return HealthCheckReq()
+        case RespKind.SetPid.value:
+            return SetPidResp(int(rest))
 
         case _:
             raise ValueError("Could not match kind for message", msg)

{cubething_psync-0.2.1.dev1 → cubething_psync-0.3.0/src/cubething_psync.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cubething_psync
-Version: 0.2.1.dev1
+Version: 0.3.0
 Summary: Simple project synchronization tool.
 Author-email: ada mandala <ada@cubething.dev>
 License-Expression: MIT

{cubething_psync-0.2.1.dev1 → cubething_psync-0.3.0}/src/cubething_psync.egg-info/SOURCES.txt

@@ -17,4 +17,5 @@ src/cubething_psync.egg-info/top_level.txt
 src/server/__init__.py
 src/server/__main__.py
 src/server/args.py
-src/server/main.py
+src/server/main.py
+test/test_integration.py

cubething_psync-0.3.0/src/server/args.py

@@ -0,0 +1,107 @@
+import argparse
+from dataclasses import dataclass, field
+from os import environ
+from pathlib import Path
+
+
+@dataclass
+class Args:
+    """
+    Server arguments.
+    """
+
+    use_base_env: bool
+    """
+    flag: ``--use-base-env, -E``
+
+    Use the current environment in addition to the values specified in the websocket request.
+    """
+    cert_path: Path = Path(
+        environ.get("PSYNC_SSL_CERT_PATH", "~/.local/share/psync/cert.pem")
+    ).expanduser()
+    """
+    environ: ``PSYNC_SSL_CERT_PATH``
+
+    Path to the SSL certificate used to authenticate this server.
+    """
+    key_path: Path = Path(
+        environ.get("PSYNC_SSL_KEY_PATH", "~/.local/share/psync/key.pem")
+    ).expanduser()
+    """
+    environ: ``PSYNC_SSL_KEY_PATH``
+
+    Path to the SSL private key.
+    """
+    host: str = environ.get("PSYNC_SERVER_IP", "0.0.0.0")
+    """
+    environ: ``PSYNC_SERVER_IP``
+
+    Host IP on which to listen for incoming connections.
+    """
+    port: str = environ.get("PSYNC_SERVER_PORT", "5000")
+    """
+    environ: ``PSYNC_SERVER_PORT``
+
+    Host port on which to listen for incoming connections.
+    """
+    origins: list[str] = field(
+        default_factory=lambda: environ.get(
+            "PSYNC_ORIGINS", "localhost 127.0.0.1"
+        ).split()
+    )
+    """
+    environ: ``PSYNC_ORIGINS``
+
+    Accepted client origins. Should match the HTTP Origin header.
+    """
+    log_level: str = environ.get("PSYNC_LOG_LEVEL", "INFO").upper()
+    """
+    environ: ``PSYNC_LOG_LEVEL``
+
+    Log level.
+    """
+    user: str | None = environ.get("PSYNC_USER", None)
+    """
+    environ: ``PSYNC_USER``
+
+    User used to execute the requested binaries.
+    """
+
+
+parser = argparse.ArgumentParser(
+    prog="psync-server",
+    usage="""\
+Server for project syncrhonization.
+
+In addition to the options below, the client is configurable through environment
+variables.
+
+SSL_CERT_PATH - Path to SSL cert
+    Default: ./cert.pem
+SSL_KEY_PATH - Path to SSL key
+    Default: ./key.pem
+PSYNC_SERVER_IP - IP address on which to listen
+    Default: 0.0.0.0
+PSYNC_SERVER_PORT - Port on which to listen
+    Default: 5000
+PSYNC_ORIGINS - Space-separated list of accepted incoming IP addresses
+    Default: "127.0.0.1 localhost"
+PSYNC_LOG - Log level
+    Default: "INFO"
+PSYNC_USER - User to run the synced executables. Try not to use root.
+    Default: None (current user)
+""",
+)
+_action = parser.add_argument(
+    "--use-base-env",
+    "-E",
+    help="Use the current environment in addition to the requested values.",
+    action="store_true",
+)
+
+
+def parse_args() -> Args:
+    args = vars(parser.parse_args())
+    return Args(
+        use_base_env=args["use_base_env"],  # pyright: ignore[reportAny]
+    )

{cubething_psync-0.2.1.dev1 → cubething_psync-0.3.0}/src/server/main.py

@@ -17,11 +17,10 @@ from typing import Callable
 from websockets import (
     ConnectionClosedError,
     ConnectionClosedOK,
-    Request,
-    Response,
     ServerConnection,
 )
 from websockets.asyncio.server import serve
+from websockets.typing import Origin
 from common.data import (
     ErrorResp,
     ExitResp,
@@ -29,13 +28,14 @@ from common.data import (
     LogResp,
     OkayResp,
     OpenReq,
+    SetPidResp,
+    HealthCheckReq,
     serialize,
     deserialize,
 )
 import logging
 from common.log import InterceptHandler
 from server.args import (
-    PSYNC_LOG,
     Args,
     parse_args,
 )
@@ -45,6 +45,10 @@ pprint = PrettyPrinter().pformat
 
 @dataclass
 class PTask:
+    """
+    Simple wrapper class for task-based process execution.
+    """
+
     task: Task[None]
     process: Process
 
@@ -52,34 +56,20 @@ class PTask:
 class PsyncServer:
     """
     The main interface for the psync websocker server.
-
-    Configuration environment variables:
-        PSYNC_HOST: IP for the server.
-            Default: "0.0.0.0"
-        PSYNC_PORT: Port for the server.
-            Default: 5000
-        PSYNC_ORIGINS: Space-separated list of allowed foreign origins.
-            Default: "localhost"
-        SSL_CERT_PATH: Path to the SSL certification file.
-            Default: "./cert.pem"
-        SSL_KEY_PATH: Path to the SSL key file.
-            Default: "./key.pem"
-
-    .. _logger: https://docs.python.org/3/library/logging.html#logging-levels
     """
 
-    __args: Args
-
-    __tasks: dict[str, PTask] = {}
-    """Active sessions. A dict of IP addresses and the running log task. IP
-    addresses _must_ match those in origins."""
+    args: Args
 
+    __tasks: dict[str, dict[int, PTask]] = {}
+    """{[host: str]: {[pid: str]: PTask} }"""
     __coroutine: Task[None] | None = None
     """The main coroutine for this server."""
 
+    __force_shutdown: bool = False
+
     def __init__(self, args: Args):
         logging.debug(pprint(args))
-        self.__args = args
+        self.args = args
 
     def __get_host(self, ws: ServerConnection) -> str:
         addrs: tuple[str, str] = ws.remote_address  # pyright: ignore[reportAny]
@@ -92,16 +82,16 @@ class PsyncServer:
         """
         ssl_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
         ssl_ctx.load_cert_chain(
-            pathlib.Path(self.__args.cert_path).expanduser(),
-            pathlib.Path(self.__args.key_path).expanduser(),
+            pathlib.Path(self.args.cert_path).expanduser(),
+            pathlib.Path(self.args.key_path).expanduser(),
         )
         logging.debug(pprint(ssl_ctx.get_ca_certs()))
         server = await serve(
             (self.__handle()),
-            self.__args.host,
-            self.__args.port,
-            process_request=self.__process_request(),
+            self.args.host,
+            int(self.args.port),
             ssl=ssl_ctx,
+            origins=list(map(lambda x: Origin(f"wss://{x}"), self.args.origins)),
         )
         self.__coroutine = asyncio.create_task(server.serve_forever())
         try:
@@ -111,17 +101,6 @@ class PsyncServer:
             logging.info(f"Got error {e}")
             pass
 
-    def __process_request(
-        self,
-    ) -> Callable[[ServerConnection, Request], Response | None]:
-        def inner(ws: ServerConnection, _req: Request) -> Response | None:
-            addrs: tuple[str, str] = ws.remote_address  # pyright: ignore[reportAny]
-            (host, _port) = addrs
-            if host not in self.__args.origins:
-                return ws.respond(400, "Client address not recognized.")
-
-        return inner
-
     async def __end_session(self, ws: ServerConnection):
         host = self.__get_host(ws)
         try:
@@ -132,10 +111,16 @@ class PsyncServer:
 
     def __mk_handle_signal(self, ws: ServerConnection):
         async def inner():
-            logging.info("Gracefully shutting down...")
-            await ws.close()
-            _ = self.__coroutine.cancel()  # pyright: ignore[reportOptionalMemberAccess]
-            asyncio.get_event_loop().stop()
+            if not self.__force_shutdown:
+                logging.info("Gracefully shutting down...")
+                self.__force_shutdown = True
+                await ws.close()
+                _ = self.__coroutine.cancel()  # pyright: ignore[reportOptionalMemberAccess]
+                asyncio.get_event_loop().stop()
+                raise SystemExit(130)
+            else:
+                logging.warning("Second Ctrl-C detected, forcing shutdown.")
+                raise SystemExit(1)
 
         return lambda: asyncio.create_task(inner())
 
@@ -163,6 +148,10 @@ class PsyncServer:
                             await self.__open(req, ws)
                         case KillReq():
                             await self.__kill(req, ws)
+                        case HealthCheckReq():
+                            logging.info("Health check OK")
+                            await ws.send(serialize(OkayResp()))
+                            await ws.close()
                         case _:
                             logging.warning(f"Got unknown request {req}")
             except ConnectionClosedOK:
@@ -176,23 +165,15 @@ class PsyncServer:
 
     async def __open(self, req: OpenReq, ws: ServerConnection):
         host = self.__get_host(ws)
-        if self.__tasks.get(host) is not None:
-            ptask = self.__tasks[host]
-            logging.warn(
-                f"Cancelling previous task for host {host} (PID {ptask.process.pid})"
-            )
-            await ptask.process.kill()
-            await ptask.task.cancel()
-            self.__tasks.pop(host)
         path = pathlib.Path.expanduser(req.path).resolve()
         args = [str(path), *req.args]
-        env = req.env if not self.__args.use_base_env else {**environ, **req.env}
+        env = req.env if not self.args.use_base_env else {**environ, **req.env}
 
         info_log = f"Running `[...]/{basename(args[0])} {' '.join(args[1:])}`..."
-        if env != {} and env is not None:
+        if env != {}:
             info_log += f"\n... with env {pprint(env)}"
-        if self.__args.user is not None:
-            info_log += f"... as user {self.__args.user}"
+        if self.args.user is not None:
+            info_log += f"... as user {self.args.user}"
 
         logging.info(info_log)
 
@@ -202,14 +183,21 @@ class PsyncServer:
                 env=env,
                 stdout=asyncio.subprocess.PIPE,
                 stderr=asyncio.subprocess.STDOUT,
-                user=self.__args.user,
+                user=self.args.user,
             )
-            self.__tasks[host] = asyncio.create_task(self.__log(ws, p))
-            resp = OkayResp()
+            task = PTask(asyncio.create_task(self.__log(ws, p)), p)
+
+            if self.__tasks.get(host) is None:
+                self.__tasks[host] = {p.pid: task}
+            else:
+                self.__tasks[host][p.pid] = task
+
+            resp = SetPidResp(pid=p.pid)
             await ws.send(serialize(resp))
+
         except Exception as e:
-            logging.error(f"Failed to run process with error {e}")
-            resp = ErrorResp(msg=f"Server error: {e}")
+            logging.error(f"Failed to start process `{args[0]}` with error {e}")
+            resp = ErrorResp(f"Server error: {e}")
             await ws.send(serialize(resp))
 
     async def __log(self, ws: ServerConnection, process: asyncio.subprocess.Process):
@@ -232,33 +220,41 @@ class PsyncServer:
         except (KeyError, ConnectionClosedError, ConnectionClosedOK):
             pass
 
-    async def __kill(self, _req: KillReq, ws: ServerConnection):
+    async def __kill(self, req: KillReq, ws: ServerConnection):
         host = self.__get_host(ws)
-        ptask = self.__tasks.get(host)
-        process = ptask.process
-        task = ptask.task
-        if process is not None and task is not None:
-            logging.info(f"Killing PID {process.pid}")
-            process.kill()
-            code = await process.wait()
-            resp = ExitResp(str(code))
-            await ws.send(serialize(resp))
-        else:
-            logging.error(
-                f"Tried to kill process for host {host}, but no process was running."
-            )
-            await ws.send(
-                serialize(
-                    ErrorResp(msg="Tried to kill process, but no process was running.")
-                )
-            )
+        tasks = self.__tasks.get(host)
+        if tasks is None:
+            msg = f"Tried to kill process for host {host}, but no process was running."
+            logging.error(msg)
+            await ws.send(serialize(ErrorResp(msg)))
+            return
+
+        task = tasks.get(req.pid)
+        if task is None:
+            msg = f"Tried to kill process {req.pid}, but it was not found."
+            logging.error(msg)
+            await ws.send(serialize(ErrorResp(msg)))
+            return
+
+        process = task.process
+        logging.info(f"Killing PID {process.pid}")
+        process.kill()
+        code = await process.wait()
+        resp = ExitResp(str(code))
+        await ws.send(serialize(resp))
         await ws.close()
 
 
-def main():
+def main(args: Args | None = None):
     """Run the server as an executable."""
-    logging.basicConfig(handlers=[InterceptHandler()], level=PSYNC_LOG, force=True)
-    asyncio.run(PsyncServer(parse_args()).serve())
+    args = parse_args() if args is None else args
+    logging.basicConfig(handlers=[InterceptHandler()], level=args.log_level, force=True)
+    try:
+        asyncio.run(PsyncServer(args).serve())
+    except SystemExit as e:
+        exit(e.code)
+    except Exception:
+        exit(1)
 
 
 if __name__ == "__main__":

cubething_psync-0.3.0/test/test_integration.py

@@ -0,0 +1,83 @@
+import asyncio
+import multiprocessing
+import os
+from os.path import basename
+import re
+import sys
+from signal import Signals
+
+from testcontainers.core.generic import DockerContainer  # pyright: ignore[reportMissingTypeStubs]
+
+from client.args import Args as ClientArgs
+from client.main import PsyncClient
+from client.main import __rsync as rsync  # pyright: ignore[reportPrivateUsage]
+from test.conftest import assets_path
+
+
+def template(args: ClientArgs, server: DockerContainer, kill: bool = False):
+    try:
+        rsync(args)
+        exec_result = server.exec(["ls", str(args.destination_path())])
+        print(f"ls {args.destination_path()}\n --- \n {exec_result.output}")
+        assert exec_result.output.decode().__contains__(basename(args.target_path))
+        client = PsyncClient(args)
+
+        def run(code: int):
+            try:
+                asyncio.run(client.run())
+            except SystemExit as e:
+                assert str(e.code) == str(code)
+
+        if not kill:
+            run(0)
+        else:
+            p = multiprocessing.Process(target=run, args=[130])
+            p.start()
+            while p.pid is None:
+                pass
+            asyncio.run(asyncio.sleep(1))
+            os.kill(p.pid, Signals.SIGINT)
+            p.join(3)
+
+        # check that the pid closed
+        stdout, stderr = server.get_logs()
+        pat = re.compile(r"Running process with PID (\d+)")
+        res = pat.search(stderr.decode())
+        if res is None:
+            raise Exception("Could not get PID from stdout!")
+        pid = res.group(1)
+
+        exec_result = server.exec(
+            ["sh", "-c", f"ps -p {pid} > /dev/null; echo $?"],
+        )
+        assert exec_result.output.decode().strip() == "1"
+
+    except Exception as e:
+        print(f"Got exception:\n {e}", file=sys.stderr)
+        stdout, stderr = server.get_logs()
+        print(
+            f"Server logs:\n--- stdout ---\n{stdout.decode()}\n--- stderr ---\n{stderr.decode()}",
+            file=sys.stderr,
+        )
+        assert False
+
+
+def get_test_args(file: str, server: DockerContainer):
+    return ClientArgs(
+        target_path=assets_path.joinpath(file).__str__(),
+        ssh_args=f"-i {(assets_path / 'ssh-key').resolve()} -l psync -o StrictHostKeyChecking=no",
+        ssl_cert_path=(assets_path / "cert.pem").resolve().__str__(),
+        server_ip="127.0.0.1",
+        server_port=server.get_exposed_port(5000),
+        server_ssh_port=server.get_exposed_port(5022),
+    )
+
+
+def test_basic(server: DockerContainer):
+    args = get_test_args("example_basic.py", server)
+    template(args, server)
+
+
+def test_sigint(server: DockerContainer):
+    args = get_test_args("example.py", server)
+    template(args, server, True)

cubething_psync-0.2.1.dev1/src/client/args.py

@@ -1,94 +0,0 @@
-import argparse
-from dataclasses import dataclass
-import logging
-import os
-from pathlib import Path
-import shlex
-from common.data import deserialize_env
-
-
-@dataclass
-class Args:
-    target_path: str
-    extra: list[str]
-    env: dict[str, str]
-    args: list[str]
-
-
-parser = argparse.ArgumentParser(
-    prog="psync-client",
-    usage="""\
-Client for the psync server.
-
-In addition to the options below, the client is configurable through environment
-variables.
-
-Variable          | Default
-------------------+-------------------------------
-PSYNC_SERVER_IP   | 127.0.0.1
-PSYNC_SERVER_PORT | 5000
-PSYNC_SSH_PORT    | 5022
-PSYNC_SERVER_DEST | /home/psync/
-PSYNC_SSH_USER    | psync
-PSYNC_CERT_PATH   | ~/.local/share/psync/cert.pem
-""",
-)
-_action = parser.add_argument(
-    "--path",
-    "-p",
-    required=True,
-    help="Path to the target exectuable.",
-)
-_action = parser.add_argument(
-    "--extra",
-    "-E",
-    nargs="+",
-    help="Extra files or directories to be synced to the destination path.",
-)
-_action = parser.add_argument(
-    "--env",
-    "-e",
-    help="Environment variables to set in the remote execution environment. Variables must be space-sepated or double-quoted.",
-)
-_action = parser.add_argument(
-    "--args", "-a", help="Arguments with which to run the remote executable."
-)
-
-SERVER_IP: str = os.environ.get("PSYNC_SERVER_IP", "127.0.0.1")
-SERVER_PORT: int = int(os.environ.get("PSYNC_SERVER_PORT", "5000"))
-SERVER_SSH_PORT: int = int(os.environ.get("PSYNC_SSH_PORT", "5022"))
-SERVER_DEST: str = os.environ.get("PSYNC_SERVER_DEST", "/home/psync")
-USER: str = os.environ.get("PSYNC_SSH_USER", "psync")
-SSL_CERT_PATH: str = os.environ.get("PSYNC_CERT_PATH", "~/.local/share/psync/cert.pem")
-
-
-def parse_args() -> Args:
-    args = vars(parser.parse_args())
-
-    target_path = str(args.get("path"))
-    target_path = Path(target_path)
-    if not target_path.is_file():
-        logging.error(f"Could not file at {target_path}")
-        exit(1)
-
-    extra: list[str] = []
-    extra_raw = args.get("extra")
-    if extra_raw is not None:
-        extra = extra_raw  # pyright: ignore[reportAny]
-
-    client_args: list[str] = []
-    raw_args = args.get("args")
-    if raw_args is not None:
-        client_args = shlex.split(str(raw_args))  # pyright: ignore[reportAny]
-
-    env: dict[str, str] = dict()
-    raw_env = args.get("env")
-    if raw_env is not None:
-        env = deserialize_env(f"env='{raw_env}'")
-
-    return Args(
-        target_path=str(target_path),
-        extra=extra or [],
-        env=env,
-        args=client_args,
-    )

cubething_psync-0.2.1.dev1/src/server/args.py

@@ -1,69 +0,0 @@
-import argparse
-from dataclasses import dataclass
-from os import environ
-from pathlib import Path
-
-SSL_CERT_PATH: str = environ.get("SSL_CERT_PATH", "~/.local/share/psync/cert.pem")
-SSL_KEY_PATH: str = environ.get("SSL_KEY_PATH", "~/.local/share/psync/key.pem")
-PSYNC_HOST: str = environ.get("PSYNC_SERVER_IP", "0.0.0.0")
-PSYNC_PORT: str = environ.get("PSYNC_SERVER_PORT", "5000")
-PSYNC_ORIGINS: str = environ.get("PSYNC_ORIGINS", "localhost 127.0.0.1")
-PSYNC_LOG: str = environ.get("PSYNC_LOG", "INFO").upper()
-PSYNC_USER: str | None = environ.get("PSYNC_USER", None)
-
-
-@dataclass
-class Args:
-    use_base_env: bool
-    host: str
-    port: str
-    origins: list[str]
-    user: str | None
-    cert_path: Path
-    key_path: Path
-
-
-parser = argparse.ArgumentParser(
-    prog="psync-server",
-    usage="""\
-Server for project syncrhonization.
-
-In addition to the options below, the client is configurable through environment
-variables.
-
-SSL_CERT_PATH - Path to SSL cert
-    Default: ./cert.pem
-SSL_KEY_PATH - Path to SSL key
-    Default: ./key.pem
-PSYNC_SERVER_IP - IP address on which to listen
-    Default: 0.0.0.0
-PSYNC_SERVER_PORT - Port on which to listen
-    Default: 5000
-PSYNC_ORIGINS - Space-separated list of accepted incoming IP addresses
-    Default: "127.0.0.1 localhost"
-PSYNC_LOG - Log level
-    Default: "INFO"
-PSYNC_USER - User to run the synced executables. Try not to use root.
-    Default: None (current user)
-""",
-)
-_action = parser.add_argument(
-    "--use-base-env",
-    "-E",
-    help="Use the current environment in addition to the requested values.",
-    action="store_true",
-)
-
-
-def parse_args() -> Args:
-    args = vars(parser.parse_args())
-    return Args(
-        use_base_env=args["use_base_env"],
-        host=PSYNC_HOST,
-        port=PSYNC_PORT,
-        origins=PSYNC_ORIGINS.split(),
-        user=PSYNC_USER,
-        cert_path=Path(SSL_CERT_PATH).expanduser(),
-        key_path=Path(SSL_KEY_PATH).expanduser(),
-    )
-    print(args)