diff --git a/scripts/discovery.py b/scripts/discovery.py index 58748ac..2f37d33 100755 --- a/scripts/discovery.py +++ b/scripts/discovery.py @@ -115,7 +115,7 @@ def main(): f"Hook {hook} does not contain a valid on_discovery_updated method." ) continue - on_discovery_updated() + on_discovery_updated(robots=robots, selected_robots=selected_robots) elif hook.endswith(".bash") or hook.endswith(".sh"): executable = "bash" if hook.endswith(".bash") else "sh" subprocess.run([executable, hook], cwd=get_workspace_root()) diff --git a/scripts/hooks/discovery/20.add_routes.py b/scripts/hooks/discovery/20.add_routes.py new file mode 100644 index 0000000..c1f6d16 --- /dev/null +++ b/scripts/hooks/discovery/20.add_routes.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +from pyroute2 import IPRoute +import subprocess + + +def on_discovery_updated(robots, selected_robots, **_): + ip = IPRoute() + + routes_to_add = "" + + for robot_name, robot in robots.items(): + if robot_name not in selected_robots: + continue + + for pc_name in robot.remote_pcs: + pc = robot.remote_pcs[pc_name] + if not pc.address: + print( + f"Robot {pc.name} has no address configured for its , skipping pc" + ) + continue + else: + pc_ip = pc.address + netmask = pc.netmask + break + + if not pc_ip: + print( + f"Robot {robot_name} has no address configured for its pcs, skipping robot" + ) + continue + parts = pc_ip.split(".")[:-1] + net_ip = ".".join(parts + ["0"]) + host_ip = ".".join(parts[:2] + ["0", parts[2]]) + + existing_routes = ip.get_routes(dst=net_ip + "/" + str(netmask)) + + valid_route = False + if existing_routes: + gateway = next( + ( + value + for key, value in existing_routes[0]["attrs"] + if key == "RTA_GATEWAY" + ), + None, + ) + if gateway == host_ip: + valid_route = True + + if not valid_route: + print(f"Adding route to {net_ip}/{netmask} via host pc") + command = ( + f"ip.route('add', dst='{net_ip}/{netmask}', gateway='{host_ip}')\n" + ) + routes_to_add += command + + if routes_to_add != "": + ip_route_code = f"""from pyroute2 import IPRoute +ip = IPRoute() +{routes_to_add} +""" + print( + "Privilege escalation required to add routes, running commands with sudo..." + ) + subprocess.run(["sudo", "python3", "-c", ip_route_code], check=True) diff --git a/scripts/hooks/discovery/90.restart_daemon.py b/scripts/hooks/discovery/90.restart_daemon.py index 153a303..d0556a9 100644 --- a/scripts/hooks/discovery/90.restart_daemon.py +++ b/scripts/hooks/discovery/90.restart_daemon.py @@ -3,7 +3,7 @@ from tuda_workspace_scripts.print import * -def on_discovery_updated(): +def on_discovery_updated(**_): print_header("Restarting ROS2 daemon") if is_daemon_running(args=[]): if not shutdown_daemon(args=[], timeout=10): diff --git a/tuda_workspace_scripts/robots.py b/tuda_workspace_scripts/robots.py index 2881f7d..28b66b6 100644 --- a/tuda_workspace_scripts/robots.py +++ b/tuda_workspace_scripts/robots.py @@ -37,12 +37,16 @@ class RemotePC: def __init__( self, name: str, + address: str, hostname: str, user: str, commands: list[Command], port: int = 22, + netmask: int = 24, ): self.name = name + self.address = address + self.netmask = netmask self.hostname = hostname self.user = user self.port = port @@ -166,6 +170,8 @@ def _load_pc_from_yaml( raise ValueError(f"User not specified for remote PC {pc_name}") user = config["user"] hostname = config["hostname"] if "hostname" in config else pc_name + address = config["address"] if "address" in config else None + netmask = config["netmask"] if "netmask" in config else 24 port = config["port"] if "port" in config else 22 commands = dict(shared_commands) if "commands" in config: @@ -175,7 +181,13 @@ def _load_pc_from_yaml( continue commands[name] = _load_command_from_yaml(name, config["commands"][name]) return RemotePC( - pc_name, hostname, user, port=port, commands=list(commands.values()) + pc_name, + address, + hostname, + user, + port=port, + netmask=netmask, + commands=list(commands.values()), )