update hooks and simln plugin

Author: mplsgrant

resources/plugins/simln/charts/simln/values.yaml

@@ -4,5 +4,5 @@ image:
   tag: "d8c165d"
   pullPolicy: IfNotPresent
 command: ["sh", "-c"]
-args: ["while true; do sleep 3600; done"]
+args: ["cd /config; sim-cli"]
 defaultDataDir: /app/data

resources/plugins/simln/simln.py

@@ -1,15 +1,229 @@
-from hooks_api import post_status, pre_status
+import json
+import logging
+import random
+from pathlib import Path
+from subprocess import run
+from time import sleep
 
+from warnet.hooks import _get_plugin_directory as get_plugin_directory
+from warnet.k8s import get_pods_with_label
+from warnet.process import run_command
+from warnet.status import _get_tank_status as network_status
 
-@pre_status
-def print_something_first():
-    print("The simln plugin is enabled.")
+log = logging.getLogger("simln")
+log.setLevel(logging.DEBUG)
+console_handler = logging.StreamHandler()
+console_handler.setLevel(logging.DEBUG)
+formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+console_handler.setFormatter(formatter)
+log.addHandler(console_handler)
 
+lightning_selector = "mission=lightning"
 
-@post_status
-def print_something_afterwards():
-    print("The simln plugin executes after `status` has run.")
 
+# @post_deploy
+def deploy_helm():
+    print("SimLN Plugin ⚡")
+    plugin_dir = get_plugin_directory()
+    print(f"DIR: {plugin_dir}")
+    command = f"helm upgrade --install simln {plugin_dir}/simln/charts/simln"
+    helm_result = run_command(command)
+    print(helm_result)
 
-def run():
-    print("Running the simln plugin")
+
+def run_simln():
+    init_network()
+    fund_wallets()
+    wait_for_predicate(everyone_has_a_host)
+    log.info(warnet("bitcoin rpc tank-0000 -generate 7"))
+    # warnet("ln open-all-channels")
+    manual_open_channels()
+    log.info(warnet("bitcoin rpc tank-0000 -generate 7"))
+    wait_for_gossip_sync(2)
+    log.info("done waiting")
+    pods = get_pods_with_label(lightning_selector)
+    pod_a = pods[0].metadata.name
+    pod_b = pods[1].metadata.name
+    sample_activity = [
+        {"source": pod_a, "destination": pod_b, "interval_secs": 1, "amount_msat": 2000}
+    ]
+    log.info(f"Activity: {sample_activity}")
+    generate_activity(sample_activity)
+    log.info("Sent command. Done.")
+
+
+def generate_activity(activity: list[dict]):
+    random_digits = "".join(random.choices("0123456789", k=10))
+    plugin_dir = get_plugin_directory()
+    generate_nodes_file(activity, plugin_dir / Path("simln/charts/simln/files/sim.json"))
+    command = f"helm upgrade --install simln-{random_digits} {plugin_dir}/simln/charts/simln"
+    log.info(f"generate activity: {command}")
+    run_command(command)
+
+
+def init_network():
+    log.info("Initializing network")
+    wait_for_all_tanks_status(target="running")
+
+    warnet("bitcoin rpc tank-0000 createwallet miner")
+    warnet("bitcoin rpc tank-0000 -generate 110")
+    wait_for_predicate(lambda: int(warnet("bitcoin rpc tank-0000 getblockcount")) > 100)
+
+    def wait_for_all_ln_rpc():
+        lns = get_pods_with_label(lightning_selector)
+        for v1_pod in lns:
+            ln = v1_pod.metadata.name
+            try:
+                warnet(f"ln rpc {ln} getinfo")
+            except Exception:
+                log.info(f"LN node {ln} not ready for rpc yet")
+                return False
+        return True
+
+    wait_for_predicate(wait_for_all_ln_rpc)
+
+
+def fund_wallets():
+    log.info("Funding wallets")
+    outputs = ""
+    lns = get_pods_with_label(lightning_selector)
+    for v1_pod in lns:
+        lnd = v1_pod.metadata.name
+        addr = json.loads(warnet(f"ln rpc {lnd} newaddress p2wkh"))["address"]
+        outputs += f',"{addr}":10'
+    # trim first comma
+    outputs = outputs[1:]
+    log.info(warnet("bitcoin rpc tank-0000 sendmany '' '{" + outputs + "}'"))
+    log.info(warnet("bitcoin rpc tank-0000 -generate 1"))
+
+
+def everyone_has_a_host() -> bool:
+    pods = get_pods_with_label(lightning_selector)
+    host_havers = 0
+    for pod in pods:
+        name = pod.metadata.name
+        result = warnet(f"ln host {name}")
+        if len(result) > 1:
+            host_havers += 1
+    return host_havers == len(pods)
+
+
+def wait_for_predicate(predicate, timeout=5 * 60, interval=5):
+    log.info(
+        f"Waiting for predicate ({predicate.__name__}) with timeout {timeout}s and interval {interval}s"
+    )
+    while timeout > 0:
+        try:
+            if predicate():
+                return
+        except Exception:
+            pass
+        sleep(interval)
+        timeout -= interval
+    import inspect
+
+    raise Exception(
+        f"Timed out waiting for Truth from predicate: {inspect.getsource(predicate).strip()}"
+    )
+
+
+def wait_for_all_tanks_status(target="running", timeout=20 * 60, interval=5):
+    """Poll the warnet server for container status
+    Block until all tanks are running
+    """
+
+    def check_status():
+        tanks = network_status()
+        stats = {"total": 0}
+        # "Probably" means all tanks are stopped and deleted
+        if len(tanks) == 0:
+            return True
+        for tank in tanks:
+            status = tank["status"]
+            stats["total"] += 1
+            stats[status] = stats.get(status, 0) + 1
+        log.info(f"Waiting for all tanks to reach '{target}': {stats}")
+        return target in stats and stats[target] == stats["total"]
+
+    wait_for_predicate(check_status, timeout, interval)
+
+
+def wait_for_gossip_sync(expected):
+    log.info(f"Waiting for sync (expecting {expected})...")
+    current = 0
+    while current < expected:
+        current = 0
+        pods = get_pods_with_label(lightning_selector)
+        for v1_pod in pods:
+            node = v1_pod.metadata.name
+            chs = json.loads(run_command(f"warnet ln rpc {node} describegraph"))["edges"]
+            log.info(f"{node}: {len(chs)} channels")
+            current += len(chs)
+        sleep(1)
+    log.info("Synced")
+
+
+def warnet(cmd):
+    log.info(f"Executing warnet command: {cmd}")
+    command = ["warnet"] + cmd.split()
+    proc = run(command, capture_output=True)
+    if proc.stderr:
+        raise Exception(proc.stderr.decode().strip())
+    return proc.stdout.decode()
+
+
+def generate_nodes_file(activity, output_file: Path = Path("nodes.json")):
+    nodes = []
+
+    for i in get_pods_with_label(lightning_selector):
+        name = i.metadata.name
+        node = {
+            "id": name,
+            "address": f"https://{name}:10009",
+            "macaroon": "/config/admin.macaroon",
+            "cert": "/config/tls.cert",
+        }
+        nodes.append(node)
+
+    data = {"nodes": nodes, "activity": activity}
+
+    with open(output_file, "w") as f:
+        json.dump(data, f, indent=2)
+
+
+def manual_open_channels():
+    def wait_for_two_txs():
+        wait_for_predicate(
+            lambda: json.loads(warnet("bitcoin rpc tank-0000 getmempoolinfo"))["size"] == 2
+        )
+
+    # 0 -> 1 -> 2
+    pk1 = warnet("ln pubkey tank-0001-ln")
+    pk2 = warnet("ln pubkey tank-0002-ln")
+    log.info(f"pk1: {pk1}")
+    log.info(f"pk2: {pk2}")
+
+    host1 = ""
+    host2 = ""
+
+    while not host1 or not host2:
+        if not host1:
+            host1 = warnet("ln host tank-0001-ln")
+        if not host2:
+            host2 = warnet("ln host tank-0002-ln")
+        sleep(1)
+
+    print(
+        warnet(
+            f"ln rpc tank-0000-ln openchannel --node_key {pk1} --local_amt 100000 --connect {host1}"
+        )
+    )
+    print(
+        warnet(
+            f"ln rpc tank-0001-ln openchannel --node_key {pk2} --local_amt 100000 --connect {host2}"
+        )
+    )
+
+    wait_for_two_txs()
+
+    warnet("bitcoin rpc tank-0000 -generate 10")

src/warnet/hooks.py

@@ -36,10 +36,19 @@ def plugin():
     pass
 
 
+@click.group(name="util")
+def util():
+    """Plugin utility functions"""
+    pass
+
+
+plugin.add_command(util)
+
+
 @plugin.command()
 def ls():
-    """List all available plugins and whether they are activated."""
-    plugin_dir = get_plugin_directory()
+    """List all available plugins and whether they are activated"""
+    plugin_dir = _get_plugin_directory()
 
     if not plugin_dir:
         click.secho("Could not determine the plugin directory location.")
@@ -57,8 +66,8 @@ def ls():
 @plugin.command()
 @click.argument("plugin", type=str, default="")
 def toggle(plugin: str):
-    """Turn a plugin on or off"""
-    plugin_dir = get_plugin_directory()
+    """Toggle a plugin on or off"""
+    plugin_dir = _get_plugin_directory()
 
     if plugin == "":
         plugin_list = get_plugins_with_status(plugin_dir)
@@ -67,15 +76,19 @@ def toggle(plugin: str):
         ]
 
         plugins_tag = "plugins"
-        q = [
-            inquirer.List(
-                name=plugins_tag,
-                message="Toggle a plugin, or ctrl-c to cancel",
-                choices=formatted_list,
-            )
-        ]
-        selected = inquirer.prompt(q, theme=GreenPassion())
-        plugin = selected[plugins_tag].split("|")[0].strip()
+        try:
+            q = [
+                inquirer.List(
+                    name=plugins_tag,
+                    message="Toggle a plugin, or ctrl-c to cancel",
+                    choices=formatted_list,
+                )
+            ]
+            selected = inquirer.prompt(q, theme=GreenPassion())
+            plugin = selected[plugins_tag].split("|")[0].strip()
+        except TypeError:
+            # user cancels and `selected[plugins_tag] fails with TypeError
+            sys.exit(0)
 
     plugin_settings = read_yaml(plugin_dir / Path(plugin) / "plugin.yaml")
     updated_settings = copy.deepcopy(plugin_settings)
@@ -87,6 +100,15 @@ def toggle(plugin: str):
 @click.argument("plugin", type=str)
 @click.argument("function", type=str)
 def run(plugin: str, function: str):
+    """Run a command available in a plugin"""
+    plugin_dir = _get_plugin_directory()
+    plugins = get_plugins_with_status(plugin_dir)
+    for plugin_path, status in plugins:
+        if plugin_path.stem == plugin and not status:
+            click.secho(f"The plugin '{plugin_path.stem}' is not enabled", fg="yellow")
+            click.secho("Please toggle it on to run commands.")
+            sys.exit(0)
+
     module = imported_modules.get(f"plugins.{plugin}")
     if hasattr(module, function):
         func = getattr(module, function)
@@ -193,7 +215,7 @@ def post_{hook}(func):
 def load_user_modules() -> bool:
     was_successful_load = False
 
-    plugin_dir = get_plugin_directory()
+    plugin_dir = _get_plugin_directory()
 
     if not plugin_dir or not plugin_dir.is_dir():
         return was_successful_load
@@ -243,7 +265,7 @@ def find_hooks(module_name: str, func_name: str):
     return pre_hooks, post_hooks
 
 
-def get_plugin_directory() -> Optional[Path]:
+def _get_plugin_directory() -> Optional[Path]:
     user_dir = os.getenv(WARNET_USER_DIR_ENV_VAR)
 
     plugin_dir = Path(user_dir) / PLUGINS_LABEL if user_dir else Path.cwd() / PLUGINS_LABEL
@@ -254,6 +276,11 @@ def get_plugin_directory() -> Optional[Path]:
         return None
 
 
+@util.command()
+def get_plugin_directory():
+    click.secho(_get_plugin_directory())
+
+
 def get_version(package_name: str) -> str:
     try:
         return version(package_name)

src/warnet/k8s.py

@@ -560,4 +560,4 @@ def get_pods_with_label(label_selector: str, namespace: Optional[str] = None) ->
         return v1_pods
     except client.exceptions.ApiException as e:
         print(f"Error fetching pods: {e}")
-        return []
+        return []