diff --git a/src/bt_auto_dumper/_v2/__main__.py b/src/bt_auto_dumper/_v2/__main__.py index b1a0676..f147eaf 100644 --- a/src/bt_auto_dumper/_v2/__main__.py +++ b/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,18 +136,12 @@ 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/" @@ -169,17 +149,17 @@ def send_to_autovalidator( "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()