Fixed bug on stopping the program with systemd where it ended abrup…

…tly instead of cleanly.
fernandoenzo · Jul 27, 2022 · 17b01d5 · 17b01d5
1 parent e4350a6
commit 17b01d5
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 8 deletions.
diff --git a/setup.py b/setup.py
@@ -7,7 +7,7 @@
 
 name = 'wireguard-subnets'
 
-version = '1.0.1'
+version = '1.0.2'
 
 description = "This program performs unattended addition and removal of remote subnets accessible through its `wg` interface " \
               "to a WireGuard server's routing table.\nWorks great combined with systemd."

diff --git a/wireguard_subnets/parser.py b/wireguard_subnets/parser.py
@@ -11,6 +11,7 @@ class ARGS:
     PERIOD = None
     METRIC = None
     IPS_SUBNETS = None
+    SYSTEMD = None
 
 
 class CustomArgumentFormatter(ArgumentDefaultsHelpFormatter, RawTextHelpFormatter):

diff --git a/wireguard_subnets/utils.py b/wireguard_subnets/utils.py
@@ -18,8 +18,12 @@ def header(title: str):
     return f' {hyphens}\n  {title}\n {hyphens}'
 
 
-def run_command(command):
-    return subprocess.run(command, capture_output=True, start_new_session=True, text=True)
+def run_command(command, systemd=False):
+    cmd = []
+    if systemd:
+        cmd = ['systemd-run', '--scope', '--quiet']
+    cmd.extend(command)
+    return subprocess.run(cmd, capture_output=True, start_new_session=True, text=True)
 
 
 def create_thread(func: Callable, *args: Any, **kwargs: Any) -> Future:

diff --git a/wireguard_subnets/wireguard_subnets.py b/wireguard_subnets/wireguard_subnets.py
@@ -29,27 +29,27 @@ def close(signum, frame):
 
 
 def link_up():
-    ret = run_command(['ip', 'link', 'show', ARGS.INTERFACE])
+    ret = run_command(['ip', 'link', 'show', ARGS.INTERFACE], systemd=ARGS.SYSTEMD)
     return not ret.returncode
 
 
 def ping(ip: Union[IPv4Address, IPv6Address]):
-    ret = run_command(['ping', '-W5', '-c3', '-I', ARGS.INTERFACE, str(ip)])
+    ret = run_command(['ping', '-W5', '-c3', '-I', ARGS.INTERFACE, str(ip)], systemd=ARGS.SYSTEMD)
     return not ret.returncode
 
 
 def add_subnet(subnet: Union[IPv4Network, IPv6Network]):
-    ret = run_command(['ip', 'route', 'add', str(subnet), 'dev', ARGS.INTERFACE, 'scope', 'link', 'metric', str(ARGS.METRIC)])
+    ret = run_command(['ip', 'route', 'add', str(subnet), 'dev', ARGS.INTERFACE, 'scope', 'link', 'metric', str(ARGS.METRIC)], systemd=ARGS.SYSTEMD)
     return not ret.returncode
 
 
 def remove_subnet(subnet: Union[IPv4Network, IPv6Network]):
-    ret = run_command(['ip', 'route', 'del', str(subnet), 'dev', ARGS.INTERFACE])
+    ret = run_command(['ip', 'route', 'del', str(subnet), 'dev', ARGS.INTERFACE], systemd=ARGS.SYSTEMD)
     return not ret.returncode
 
 
 def subnet_exists(subnet: Union[IPv4Network, IPv6Network]):
-    ret = run_command(['ip', 'route', 'show', str(subnet), 'dev', ARGS.INTERFACE])
+    ret = run_command(['ip', 'route', 'show', str(subnet), 'dev', ARGS.INTERFACE], systemd=ARGS.SYSTEMD)
     return not ret.returncode and bool(ret.stdout)
 
 
@@ -80,6 +80,7 @@ def main():
         print('This program must be run with root privileges.')
         sys.exit(0)
     parse_args()
+    ARGS.SYSTEMD = not run_command(['pidof', 'systemd'], systemd=False).returncode
     print(header('WireGuard Subnets'))
     print(f'Interface: {ARGS.INTERFACE}\nRecheck period: {ARGS.PERIOD}s\nMetric: {ARGS.METRIC}\nSubnets: ')
     collections.deque(print(f'  • {subnet} behind {ip}') for ip_subnets in ARGS.IPS_SUBNETS for ip, subnets in ip_subnets.items() for subnet in subnets)