feat: add devices management with config generator (#4)

Author: msterhuj

labctl/commands/__init__.py

@@ -1 +1,2 @@
 from .config import app as config_app
+from .devices import app as devices_app

labctl/commands/config.py

@@ -15,7 +15,7 @@ def show():
     config = Config()
     api_token = config.api_token
     if api_token:
-        me = APIDriver().get("/me")
+        me = APIDriver().get("/me").json()
         # todo handle old token and valid token
         if me.get("email"):
             api_token = "Logged in as " + me["email"]
@@ -27,6 +27,7 @@ def show():
     table.add_column("Key", style="cyan")
     table.add_column("Value", style="magenta")
     table.add_row("API URL", config.api_endpoint)
+    table.add_row("API User", config.username)
     table.add_row("API Token", api_token)
     console.print(table)
 

labctl/commands/devices.py

@@ -0,0 +1,71 @@
+from os import getcwd
+
+import typer
+
+from rich.table import Table
+
+from labctl.core import Config, APIDriver, console
+from labctl.core import cli_ready, wireguard
+
+app = typer.Typer()
+
+@app.command(name="list")
+@cli_ready
+def list_devices():
+    """
+    List devices
+    """
+    config = Config()
+    devices = APIDriver().get("/devices/" + config.username).json()
+    table = Table(title=":computer: Devices")
+    table.add_column("ID", style="cyan")
+    table.add_column("Name", style="magenta")
+    table.add_column("IPv4", style="green")
+    table.add_column("RX Bytes", style="blue")
+    table.add_column("TX Bytes", style="yellow")
+    table.add_column("Remote IP", style="red")
+
+    for device in devices:
+        table.add_row(
+            device["id"],
+            device["name"],
+            device["ipv4"],
+            device["rx_bytes"],
+            device["tx_bytes"],
+            device["remote_ip"],
+        )
+    console.print(table)
+
+@app.command(name="create")
+@cli_ready
+def create_device(name: str = typer.Argument(..., help="The device name")):
+    """
+    Create a device
+    """
+    rsp = APIDriver().post(
+        f"/devices/{Config().username}",
+        json={"name": name},
+        additional_headers={"Content-Type": "application/json"}
+    )
+    if rsp.status_code >= 200 < 300:
+        console.print(f"Device {name} created :tada:")
+        data = rsp.json()
+        config_path = f"/{getcwd()}/{name}.conf"
+        wireguard.generate_config(data["device"], data["private_key"], config_path)
+        console.print(f"Configuration file saved to {config_path}")
+        return
+    console.print(f"Error creating device {name} ({rsp.status_code})")
+
+@app.command(name="delete")
+@cli_ready
+def delete_device(
+    device_id: str = typer.Argument(..., help="The device ID")
+):
+    """
+    Delete a device
+    """
+    rsp = APIDriver().delete(f"/devices/{Config().username}/{device_id}")
+    if rsp.status_code == 200:
+        console.print(f"Device {device_id} deleted :fire:")
+        return
+    console.print(f"Error deleting device {device_id} ({rsp.status_code})")

labctl/core/__init__.py

@@ -2,3 +2,4 @@
 from .api import APIDriver
 from .console import console
 from .decorators import cli_ready
+from . import wireguard

labctl/core/api.py

@@ -22,12 +22,19 @@ def __init__(self):
         }
 
     def validate_token(self):
-        return self.get("/token/verify").get("valid", False)
+        return self.get("/token/verify").json().get("valid", False)
 
-    def get(self, path: str):
-        return requests.get(self.api_url + path, headers=self.headers).json()
+    def get(self, path: str) -> requests.Response:
+        return requests.get(self.api_url + path, headers=self.headers)
 
-    def post(self, path: str, data: dict = {}, additional_headers: dict = {}):
+    def post(self, path: str, data: dict = {}, json: dict = {}, additional_headers: dict = {}) -> requests.Response:
         headers = self.headers
         headers.update(additional_headers)
-        return requests.post(self.api_url + path, headers=headers, data=data).json()
+        if data:
+            return requests.post(self.api_url + path, headers=headers, data=data)
+        if json:
+            return requests.post(self.api_url + path, headers=headers, json=json)
+        return requests.post(self.api_url + path, headers=headers)
+
+    def delete(self, path: str) -> requests.Response:
+        return requests.delete(self.api_url + path, headers=self.headers)

labctl/core/config.py

@@ -8,6 +8,7 @@
 class Config:
 
     api_endpoint: str
+    username: str
     api_token: str
     token_type: str
 
@@ -53,4 +54,4 @@ def ready(self):
         """
         Check if the configuration is ready to be used
         """
-        return all([self.api_endpoint, self.api_token, self.token_type])
+        return all([self.api_endpoint, self.api_token, self.token_type, self.username])

labctl/core/wireguard.py

@@ -0,0 +1,16 @@
+import wgconfig
+
+def generate_config(device: dict, private_key: str, config_file: str):
+    wg = wgconfig.WGConfig(config_file)
+    wg.add_attr(None, "PrivateKey", private_key)
+    wg.add_attr(None, "Address", f"{device['ipv4']}/32, {device['ipv6']}/128")
+    wg.add_attr(None, "DNS", ", ".join(device["dns"]))
+    wg.add_attr(None, "MTU", device["mtu"])
+
+    wg.add_peer(device["server_public_key"], "# LaboInfra WireGuard Server Client")
+    wg.add_attr(device["server_public_key"], "Endpoint", device["endpoint"])
+    wg.add_attr(device["server_public_key"], "AllowedIPs", ", ".join(device["allowed_ips"]))
+    wg.add_attr(device["server_public_key"], "PersistentKeepalive", device["persistent_keepalive"])
+    wg.add_attr(device["server_public_key"], "PresharedKey", device["preshared_key"])
+
+    wg.write_file()

labctl/main.py

@@ -13,6 +13,7 @@
 app = typer.Typer()
 
 app.add_typer(commands.config_app, name="config", help="Manage the configuration")
+app.add_typer(commands.devices_app, name="devices", help="Manage vpn devices")
 
 @app.callback()
 def callback():
@@ -39,7 +40,7 @@ def me(
     Print the current status of the fastonboard-api account
     """
     api_driver = APIDriver()
-    data = api_driver.get("/me")
+    data = api_driver.get("/me").json()
     if json:
         print(dumps(data))
         return
@@ -66,12 +67,12 @@ def sync():
     Ask FastOnBoard-API to sync your account onto the vpn and openstack services
     """
     api = APIDriver()
-    me = api.get("/me")
-    task_id = api.get("/users/" + me['username'] + "/sync")
+    me = api.get("/me").json()
+    task_id = api.get("/users/" + me['username'] + "/sync").json()
     typer.echo(f"Syncing account for user {me['username']} this may take a while...")
     typer.echo("Task ID: " + task_id.get("id"))
     while True:
-        task = api.get("/users/" + me['username'] + "/sync/" + task_id.get("id"))
+        task = api.get("/users/" + me['username'] + "/sync/" + task_id.get("id")).json()
         if task.get("status") == "SUCCESS":
             typer.echo("Sync successful")
             break
@@ -81,11 +82,16 @@ def sync():
         sleep(1)
 
 @app.command()
-def login(username: Annotated[str, typer.Argument(help="The username to authenticate with")]):
+def login(username: str = typer.Option(None, help="The username to login with")):
     """
     Login to the FastOnBoard-API server
     Enter your password when prompted or set LABCTL_API_ENDPOINT_PASSWORD
     """
+    env_user = environ.get("LABCTL_API_ENDPOINT_USERNAME")
+    username = Config().username or username or env_user
+    if not username:
+        username = typer.prompt("Enter your username")
+
     env_pass = environ.get("LABCTL_API_ENDPOINT_PASSWORD")
     if env_pass:
         password = env_pass
@@ -103,7 +109,7 @@ def login(username: Annotated[str, typer.Argument(help="The username to authenti
         'password': password,
     }, additional_headers={
         'Content-Type': 'application/x-www-form-urlencoded',
-    })
+    }).json()
     if 'detail' in data:
         if "Method Not Allowed" in data['detail']:
             console.print("[red]Error: Invalid endpoint or path to api[/red]")
@@ -112,6 +118,7 @@ def login(username: Annotated[str, typer.Argument(help="The username to authenti
         return
     if 'access_token' in data:
         config = Config()
+        config.username=username
         config.api_token=data['access_token']
         config.token_type=data["token_type"]
         config.save()

poetry.lock

@@ -25,6 +25,7 @@ requests = "^2.32.3"
 pyyaml = "^6.0.2"
 typer = {extras = ["all"], version = "^0.12.5"}
 colorama = "^0.4.6"
+wgconfig = "^1.0.4"
 
 
 [tool.poetry.dev-dependencies]