Refactor into library

.github/workflows/build.yml

@@ -20,4 +20,4 @@ jobs:
             extra-substituters = https://cache.thymis.io
       - uses: DeterminateSystems/magic-nix-cache-action@v2
       - name: build controller
-        run: nix build .#default --print-build-logs
+        run: nix build .#default --print-build-logs

.github/workflows/test.yml

@@ -0,0 +1,26 @@
+name: "Test"
+on:
+  pull_request:
+  merge_group:
+  push:
+    branches:
+      - master
+
+jobs:
+  test-pre-commit:
+    strategy:
+      matrix:
+        python-version: ["3.12"]
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Setup Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install pre-commit
+        run: |
+          python -m pip install pre-commit
+      - name: Run pre-commit
+        run: |
+          pre-commit run --all-files

.gitignore

@@ -3,4 +3,4 @@ venv
 __pycache__
 result
 credentials.json
-.pytest_cache
+.pytest_cache

.pre-commit-config.yaml

@@ -0,0 +1,42 @@
+# See https://pre-commit.com for more information
+# See https://pre-commit.com/hooks.html for more hooks
+default_language_version:
+  python: python3.12
+
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.6.0
+    hooks:
+      - id: trailing-whitespace
+      - id: end-of-file-fixer
+      - id: check-added-large-files
+      - id: check-ast
+      - id: check-yaml
+      - id: check-json
+        exclude: "^frontend/tsconfig.json$"
+      - id: check-toml
+      - id: mixed-line-ending
+  - repo: https://github.com/psf/black
+    rev: "23.10.1"
+    hooks:
+      - id: black
+  - repo: https://github.com/pycqa/isort
+    rev: "5.13.2"
+    hooks:
+      - id: isort
+        name: isort (python)
+        args: ["--profile", "black"]
+  - repo: https://github.com/PyCQA/autoflake
+    rev: v2.3.1
+    hooks:
+      - id: autoflake
+        args:
+          - --in-place
+          - --expand-star-imports
+          - --remove-duplicate-keys
+          - --remove-unused-variables
+          - --remove-all-unused-imports
+  - repo: https://github.com/nix-community/nixpkgs-fmt
+    rev: v1.3.0
+    hooks:
+      - id: nixpkgs-fmt

flake.nix

@@ -8,27 +8,27 @@
   };
 
   outputs = { nixpkgs, poetry2nix, ... }:
-  let
-    eachSystem = nixpkgs.lib.genAttrs ["x86_64-linux" "aarch64-linux" "x86_64-darwin" "aarch64-darwin"];
-  in
-  {
-    packages = eachSystem (system:
-      let
-        pkgs = nixpkgs.legacyPackages.${system};
-        inherit (poetry2nix.lib.mkPoetry2Nix { inherit pkgs; }) mkPoetryApplication;
-      in
-      {
-        default = mkPoetryApplication {
-          name = "http-network-relay";
-          projectDir = ./.;
-          python = pkgs.python312;
-          checkPhase = ''
-            runHook preCheck
-            pytest --session-timeout=10
-            runHook postCheck
-          '';
-        };
-      }
-    );
-  };
+    let
+      eachSystem = nixpkgs.lib.genAttrs [ "x86_64-linux" "aarch64-linux" "x86_64-darwin" "aarch64-darwin" ];
+    in
+    {
+      packages = eachSystem (system:
+        let
+          pkgs = nixpkgs.legacyPackages.${system};
+          inherit (poetry2nix.lib.mkPoetry2Nix { inherit pkgs; }) mkPoetryApplication;
+        in
+        {
+          default = mkPoetryApplication {
+            name = "http-network-relay";
+            projectDir = ./.;
+            python = pkgs.python312;
+            checkPhase = ''
+              runHook preCheck
+              pytest --session-timeout=10
+              runHook postCheck
+            '';
+          };
+        }
+      );
+    };
 }

http_network_relay/access_client.py

@@ -1,22 +1,52 @@
 import argparse
 import asyncio
-
 import base64
 import os
 import sys
+from typing import Literal, Union
 
 import websockets
+from pydantic import BaseModel, Field
 from websockets.asyncio.client import connect
 
-from .pydantic_models import (
-    AccessClientToRelayMessage,
-    AtRStartMessage,
-    AtRTCPDataMessage,
-    RelayToAccessClientMessage,
-    RtAErrorMessage,
-    RtAStartOKMessage,
-    RtATCPDataMessage,
-)
+
+class AccessClientToRelayMessage(BaseModel):
+    inner: Union["AtRStartMessage", "AtRTCPDataMessage"] = Field(discriminator="kind")
+
+
+class AtRStartMessage(BaseModel):
+    kind: Literal["start"] = "start"
+    connection_target: str
+    target_ip: str
+    target_port: int
+    protocol: str
+    secret: str
+
+
+class AtRTCPDataMessage(BaseModel):
+    kind: Literal["tcp_data"] = "tcp_data"
+    data_base64: str
+
+
+class RelayToAccessClientMessage(BaseModel):
+    inner: Union["RtAErrorMessage", "RtAStartOKMessage", "RtATCPDataMessage"] = Field(
+        discriminator="kind"
+    )
+
+
+class RtAErrorMessage(BaseModel):
+    kind: Literal["error"] = "error"
+    message: str
+
+
+class RtAStartOKMessage(BaseModel):
+    kind: Literal["start_ok"] = "start_ok"
+
+
+class RtATCPDataMessage(BaseModel):
+    kind: Literal["tcp_data"] = "tcp_data"
+    data_base64: str
+
 
 parser = argparse.ArgumentParser(
     description="Connect to the HTTP network relay, "
@@ -131,5 +161,6 @@ async def read_stdin_and_send():
 def main():
     asyncio.run(async_main())
 
+
 if __name__ == "__main__":
-    main()
+    main()