renamed cleanup_services.py

README.md

@@ -17,6 +17,13 @@ Use the `install_service.py` script to install.
 > [!CAUTION]
 > `example.service` does not stop when disconnecting from `/run/example.sock`
 
+## Scripts
+
+- `install_service.py`: Installs the services on the system
+- `run_service.py`: Used to demo running the code tunnel from a python function
+- `test_socket.py`: Opens a connection to the socket
+- `cleanup_services.sh`: Stops services and removes symlinks
+
 ## Other Commands
 
 Reload systemd units

cleanup_services.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python
+
+# for service in $(systemctl list-units --type=service --no-legend | grep example | awk '{print $1}'); do
+#     sudo systemctl stop $service
+#     echo "Stopped $service"
+# done
+
+import subprocess
+from pathlib import Path
+from typing import List
+
+
+def list_units() -> List[str]:
+    result = subprocess.run(
+        ['systemctl', 'list-units', '--type=service', '--no-legend'],
+        check=True,
+        text=True,
+        capture_output=True,
+    )
+    return [line.split()[0] for line in result.stdout.splitlines()]
+
+
+def stop_service(service: str):
+    result = subprocess.run(['sudo', 'systemctl', 'stop', service])
+    if result.returncode == 0:
+        print(f'Stopped {service}')
+    return result.returncode
+
+
+def cleanup_service(service: str):
+    stop_result = stop_service(service)
+
+    service_file = Path('/etc/systemd/system') / service
+    if service_file.exists():
+        service_file.unlink()
+
+    if stop_result == 0:
+        subprocess.run(['sudo', 'systemctl', 'daemon-reload'], check=True)
+        print('Reloaded services')
+
+
+if __name__ == '__main__':
+    names = list_units()
+    for service in ['example.service', 'example.socket']:
+        cleanup_service(service)

cleanup_services.sh

@@ -1,6 +0,0 @@
-#!/bin/bash
-
-for service in $(systemctl list-units --type=service --no-legend | grep example | awk '{print $1}'); do
-    sudo systemctl stop $service
-    echo "Stopped $service"
-done

requirements.txt

@@ -0,0 +1 @@
+rich

run_service.py

@@ -1,27 +1,30 @@
-#!/usr/bin/env python3
+#!/usr/bin/env python
 
-import os
 import subprocess
 
-from rich import print
 
-os.chdir('/root')
-cmd = ['code', 'tunnel']
-# cmd = ['echo', '"Hello world"']
-
-
-print('PATH:\n' + '\n'.join(f' - {p}' for p in os.environ['PATH'].split(':')))
-try:
-    result = subprocess.run(
-        cmd,
-        shell=True,
-        env=os.environ,
-        text=True,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.STDOUT,
+def run_tunnel():
+    process = subprocess.Popen(
+        ['/root/code', 'tunnel'], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True
     )
-    result.check_returncode()
-except subprocess.CalledProcessError as e:
-    print(e)
+
+    with process.stdout:
+        try:
+            for line in iter(process.stdout.readline, ''):
+                print(line, end='')
+        except KeyboardInterrupt:
+            print('KeyboardInterrupt')
         finally:
-    print(result.stdout)
+            process.terminate()
+            try:
+                # Give some time to gracefully shutdown
+                returncode = process.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                print('Timeout expired, sending SIGKILL')
+                process.kill()
+                returncode = process.wait()
+    print(f'Process finished with {returncode}')
+
+
+if __name__ == '__main__':
+    run_tunnel()

test_socket.py

@@ -9,16 +9,10 @@ console = Console()
 print = console.print
 
 with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock:
-    try:
     sock.connect('/run/example.sock')
-    except Exception:
-        raise
-    else:
     with console.status('Connected...'):
         while True:
             try:
                 sleep(0.1)
             except KeyboardInterrupt:
                 break
-    finally:
-        print('Exiting')