feat: add self-commands request to server

src/bt_auto_dumper/_v2/__main__.py

@@ -1,9 +1,9 @@
 import argparse
 import configparser
 import json
+import logging
 import os
 import pathlib
-import re
 import subprocess
 import time
 import zipfile
@@ -20,22 +20,22 @@ def main(apiver: str | None = None):
     apiver = apiver or pathlib.Path(__file__).parent.name
     parser = argparse.ArgumentParser(description=f"BT Auto Dumper CLI {apiver}")
     parser.add_argument("--note", help="Comment or note for the operation", type=str, default="")
-    parser.add_argument("subnet_identifier", help="Subnet Identifier", type=str)
-    parser.add_argument("autovalidator_address", help="AutoValidator Address", type=str)
+    parser.add_argument("--subnet_identifier", help="Subnet Identifier", type=str, default="")
+    parser.add_argument("--autovalidator_address", help="AutoValidator Address", type=str, default="")
     parser.add_argument(
-        "subnet_realm",
-        help="Subnet Realm",
+        "--chain",
+        help="Specify the chain to use",
         type=str,
         choices=["testnet", "mainnet", "devnet"],
         default="mainnet",
     )
-    parser.add_argument("--set-autovalidator-address", help="Set a new autovalidator address", type=str)
-    parser.add_argument("--set-codename", help="Set a new Subnet Identifier codename", type=str)
+    parser.add_argument("--set-autovalidator-address", help="Set a new autovalidator address", type=str, default="")
+    parser.add_argument("--set-codename", help="Set a new Subnet Identifier codename", type=str, default="")
 
     args = parser.parse_args()
 
     # Get configuration directory from env variable.
-    config_base_dir = os.getenv("CONFIG_DIR")
+    config_base_dir = os.getenv("CONFIG_DIR", default="~/.config/bt-auto-dumper")
 
     # Check if the CONFIG_DIR environment variable is set
     if not config_base_dir:
@@ -53,86 +53,72 @@ def main(apiver: str | None = None):
             new_autovalidator_address=args.set_autovalidator_address,
             new_codename=args.set_codename,
         )
-        print(f"Configuration updated successfully at {config_path}")
+        logging.info(f"Configuration updated successfully at {config_path}")
 
     if not (subnet_identifier := args.subnet_identifier) or not (autovalidator_address := args.autovalidator_address):
-        autovalidator_address, subnet_identifier = load_config(config_path=config_path)
-    dump_and_upload(subnet_identifier, args.subnet_realm, autovalidator_address, args.note)
+        autovalidator_address, subnet_identifier, __, __, __ = load_config(config_path=config_path)
+    __, __, wallet_name, wallet_hotkey, wallet_path = load_config(config_path=config_path)
+    wallet = bt.wallet(name=wallet_name, hotkey=wallet_hotkey, path=wallet_path)
+    dump_and_upload(subnet_identifier, args.chain, wallet, autovalidator_address, args.note)
 
 
-def dump_and_upload(subnet_identifier: str, subnet_realm: str, autovalidator_address: str, note: str):
+def dump_and_upload(
+    subnet_identifier: str, subnet_chain: str, wallet: bt.wallet, autovalidator_address: str, note: str
+):
     """
-    Dump and upload the output of the commands to the AutoValidator
-    Args:
-        subnet_identifier: Subnet Identifier
-        subnet_realm: Subnet Realm
-        autovalidator_address: AutoValidator Address
-        note: Comment or note for the operation
+    Dump and upload the logs of the commands to the AutoValidator
     Example:
         dump_and_upload("computehorde", "mainnet", "http://localhost:8000", "Test")
     """
-    subnets = {
-        "computehorde": ["echo 'Mainnet Command 1'", "echo 'Mainnet Command 2'"],
-        "omron": ["echo 'Mainnet Command 1'", "echo 'Mainnet Command 2'"],
-    }
-
-    wallet = bt.wallet(name="validator", hotkey="validator-hotkey")
-    normalized_subnet_identifier = re.sub(r"[_\-.]", "", str.lower(subnet_identifier))
-    commands = {}
-    if normalized_subnet_identifier in subnets:
-        commands = {normalized_subnet_identifier: subnets[normalized_subnet_identifier]}
 
+    commands = get_commands_from_server(subnet_identifier, subnet_chain, wallet, autovalidator_address)
     if not commands:
-        print(f"Subnet identifier {subnet_identifier} not found.")
+        logging.error(f"Subnet dumper commands of {subnet_identifier} not found.")
         return
+    logging.info(f"Subnet dumper commands of {subnet_identifier} retrieved successfully. {commands}")
     output_files = []
-    for subnet_id, cmds in commands.items():
-        for i, command in enumerate(cmds, start=1):
-            output_file = f"{subnet_id}_{i}.txt"
-            with open(output_file, "w") as f:
-                f.write(f"Command: {command}\n")
-                result = subprocess.run(command, shell=True, capture_output=True, text=True)
-                f.write(result.stdout)
-            output_files.append(output_file)
-
-    zip_filename = f"{normalized_subnet_identifier}-output.zip"
+    for i, command in enumerate(commands, start=1):
+        output_file = f"{subnet_identifier}_{i}.txt"
+        with open(output_file, "w") as f:
+            f.write(f"Command: {command}\n")
+            result = subprocess.run(command, shell=True, capture_output=True, text=True)
+            f.write(result.stdout)
+        output_files.append(output_file)
+
+    zip_filename = f"{subnet_identifier}-output.zip"
     with zipfile.ZipFile(zip_filename, "w") as zipf:
         for file in output_files:
             zipf.write(file)
-    send_to_autovalidator(zip_filename, wallet, autovalidator_address, note, normalized_subnet_identifier, subnet_realm)
+    send_to_autovalidator(zip_filename, wallet, autovalidator_address, note, subnet_identifier, subnet_chain)
 
 
 def make_signed_request(
-    method: str, url: str, headers: dict, file_path: str, wallet: bt.wallet, subnet_realm: str
+    method: str, url: str, headers: dict, file_path: str, wallet: bt.wallet, subnet_chain: str
 ) -> requests.Response:
     """
-    Make a signed request to the AutoValidator
-    Args:
-        method: HTTP method
-        url: URL
-        headers: HTTP headers
-        file_path: File path
-        wallet: Wallet object
-    Returns:
-        Response object
     Example:
-        make_signed_request(
+        >>> make_signed_request(
             "POST",
             "http://localhost:8000/api/v1/files/",
             {"Note": "Test"},
-            {"file": open("test.zip", "rb")},
-            wallet
+            "/path/test.zip",
+            wallet,
+            "mainnet",
         )
+
     """
     headers["Nonce"] = str(time.time())
     headers["Hotkey"] = wallet.hotkey.ss58_address
-    headers["Realm"] = subnet_realm
-    files = {"file": open(file_path, "rb")}
-    file = files.get("file")
+    headers["Realm"] = subnet_chain
     file_content = b""
-    if isinstance(file, BufferedReader):
-        file_content = file.read()
-        file.seek(0)
+    files = None
+    if file_path:
+        files = {"file": open(file_path, "rb")}
+        file = files.get("file")
+
+        if isinstance(file, BufferedReader):
+            file_content = file.read()
+            file.seek(0)
     headers_str = json.dumps(headers, sort_keys=True)
     data_to_sign = f"{method}{url}{headers_str}{file_content.decode(errors='ignore')}".encode()
     signature = wallet.hotkey.sign(
@@ -150,36 +136,30 @@ def send_to_autovalidator(
     autovalidator_address: str,
     note: str,
     subnet_identifier: str,
-    subnet_realm: str,
+    subnet_chain: str,
 ):
     """
-    Send the dump file to the AutoValidator
-    Args:
-        zip_filename: Zip file name
-        wallet: Wallet object
-        autovalidator_address: AutoValidator Address
-        note: Comment or note for the operation
-        subnet_identifier: Subnet Identifier
     Example:
-        send_to_autovalidator("test.zip", wallet, "http://localhost:8000", "Test", "computehorde")
+        >>> send_to_autovalidator("test.zip", wallet, "http://localhost:8000", "Test", "computehorde", "mainnet")
+
     """
     url = f"{autovalidator_address}/api/v1/files/"
 
     headers = {
         "Note": note,
         "SubnetID": subnet_identifier,
     }
-    response = make_signed_request("POST", url, headers, zip_filename, wallet, subnet_realm)
+    response = make_signed_request("POST", url, headers, zip_filename, wallet, subnet_chain)
     if response.status_code == 201:
-        print("File successfully uploaded and resource created.")
+        logging.info("File successfully uploaded and resource created.")
     elif response.status_code == 200:
-        print("Request succeeded.")
+        logging.warning("Request succeeded.")
     else:
-        print(f"Failed to upload file. Status code: {response.status_code}")
-        print(response.text)
+        logging.error(f"Failed to upload file. Status code: {response.status_code}")
+        logging.error(response.text)
 
 
-def load_config(config_path: str) -> tuple[str, str]:
+def load_config(config_path: str) -> tuple[str, str, str, str, str]:
     """
     Load the configuration from the config file.
 
@@ -206,10 +186,19 @@ def load_config(config_path: str) -> tuple[str, str]:
     try:
         autovalidator_address = config.get("autovalidator", "autovalidator_address")
         subnet_identifier = config.get("autovalidator", "codename")
+        bittensor_wallet_name = config.get("autovalidator", "bittensor_wallet_name", fallback="validator")
+        bittensor_wallet_hotkey = config.get("autovalidator", "bittensor_wallet_hotkey", fallback="validator-hotkey")
+        bittensor_wallet_path = config.get("autovalidator", "bittensor_wallet_path", fallback="~/.bittensor/wallets")
     except Exception as e:
         raise KeyError(f"Configuration error: Missing in the config file. \n Error:{e}")
 
-    return autovalidator_address, subnet_identifier
+    return (
+        autovalidator_address,
+        subnet_identifier,
+        bittensor_wallet_name,
+        bittensor_wallet_hotkey,
+        bittensor_wallet_path,
+    )
 
 
 def update_confg(config_path: str, new_autovalidator_address: str, new_codename: str):
@@ -250,5 +239,34 @@ def update_confg(config_path: str, new_autovalidator_address: str, new_codename:
         raise RuntimeError(f"Failed to write to the configuration file: {config_path}.\n Error: {e}")
 
 
+def get_commands_from_server(
+    subnet_identifier: str, subnet_chain: str, wallet: bt.wallet, autovalidator_address: str
+) -> list:
+    """
+    Example:
+        >>> get_commands_from_server("computehorde", "mainnet", wallet, "http://localhost:8000")
+        [
+            "ps awux",
+            "docker ps",
+            "uptime",
+            "free -m",
+        ]
+
+    """
+    url = f"{autovalidator_address}/api/v1/commands/"
+    headers = {
+        "Note": "",
+        "SubnetID": subnet_identifier,
+    }
+    response = make_signed_request("GET", url, headers, "", wallet, subnet_chain)
+    if response.status_code == 200:
+        data = response.json()
+        return data
+    else:
+        logging.error(f"Failed to get commands. Status code: {response.status_code}")
+        logging.error(response.text)
+        return []
+
+
 if __name__ == "__main__":
     main()