import argparse import json import os import pyfiglet import sys import textwrap # _ _ _ # |_|___|_| |_ # | | | | _| # |_|_|_|_|_| SCRIPT_VERSION = "0.1" class bcolors: PURPLE = '\033[95m' BLUE = '\033[94m' CYAN = '\033[96m' GREEN = '\033[92m' YELLOW = '\033[93m' RED = '\033[91m' BOLD = '\033[1m' UNDERLINE = '\033[4m' ENDC = '\033[0m' # path where the entries are stored DIR_PATH = "/home/vashqlf/Code/Python/mm/pages/" # order of the rows in the table output KEY_ORDER = ["name", "ipv4", "ipv6", "os", "desc"] # commands that need a name value NAME_CMDS = ["add", "group", "edit", "rm", "view"] TERM_SIZE = os.get_terminal_size() def init_function() -> None: if not os.path.exists(DIR_PATH): os.mkdir(DIR_PATH) # _ _ # _ _ ___|_|_ _ ___ ___ ___ ___| | # | | | | | | | -_| _|_ -| .'| | # |___|_|_|_|\_/|___|_| |___|__,|_| def page_path(hostname: str) -> str: return DIR_PATH + f"{hostname}.json" # checks if the executed command needs a "name" argument def name_required() -> bool: common = set(sys.argv) & set(NAME_CMDS) return bool(common) def parse_args() -> dict: parser = argparse.ArgumentParser(description="Small asset management tool. ", formatter_class=argparse.RawTextHelpFormatter) parser.add_argument("action", nargs='?', default=os.getcwd(), help=textwrap.dedent(""" view - Show a page add - Create a page edit - Edit a page ls - List existing pages """), ) parser.add_argument("-n", "--name", required=name_required(), help="name of the page, needs to be unique") parser.add_argument("--ipv4", required=False) parser.add_argument("--ipv6", required=False) parser.add_argument("--os", required=False) parser.add_argument("--desc", required=False) parser.add_argument("--force", required=False, action='store_true') args = parser.parse_args() return vars(args) def print_error(msg: str): print(f"{bcolors.RED}{msg}{bcolors.ENDC}") def get_groups() -> List[str]: with os.scandir(DIR_PATH) as entries: groups = [e.name for e in entries if e.is_dir()] return groups def get_standalone_pages() -> List[str]: with os.scandir(DIR_PATH) as entries: pages = [e.name for e in entries if e.is_file()] return pages def with_last(iterable): it = iter(iterable) try: prev = next(it) except StopIteration: return for item in it: yield prev, False prev = item yield prev, True def get_mm_info() -> dict: dirname = os.path.dirname(__file__) absolute_path = os.path.join(dirname, DIR_PATH) data = { "groups": {}, "pages": 0, "dirpath": absolute_path, "version": SCRIPT_VERSION } try: items = os.listdir(DIR_PATH) for item in items: if os.path.isfile(DIR_PATH + item): data["pages"] += 1 else: group_items = os.listdir(DIR_PATH + item) data["groups"][item] = len(group_items) data["pages"] += len(group_items) except FileNotFoundError: print_error("Page-path corrupted.") return data # _ _ _ # ___ _| |_| |_ ___ ___ # | -_| . | | _| . | _| # |___|___|_|_| |___|_| def add_group(name: str): if os.path.exists(DIR_PATH + name): print_error(f"Group {name} already exists.") else: os.mkdir(DIR_PATH + name) def delete_item(name: str): if not os.path.exists(DIR_PATH + name): print_error(f"Item {name} doesn't exist.") elif os.path.isfile(DIR_PATH + name): os.remove(DIR_PATH + name) else: for file in os.listdir(DIR_PATH + name): os.remove(DIR_PATH + name + "/" + file) os.rmdir(DIR_PATH + name) def edit_page(args: dict): data = {} for k in KEY_ORDER: if k in args and args[k]: data[k] = args[k] try: page = {} with open(page_path(data["name"]), "r") as fp: page = json.load(fp) page.update(data) with open(page_path(data["name"]), "w") as fp: json.dump(page, fp, indent=4) except FileNotFoundError: print_error(f"Page {data['name']} doesn't exist.") except TypeError: print_error("Values must be strings, integers or floats.") def add_page(args: dict): data = {} for k in KEY_ORDER: if k in args and args[k]: data[k] = args[k] try: if os.path.exists(page_path(data["name"])) and not args["force"]: print_error(f"Page {data['name']} already exists, use --force to overwrite.") exit() with open(page_path(data["name"]), "w") as fp: json.dump(data, fp, indent=4) except FileNotFoundError: print_error("Page-Path corrupted, restart mm.") except TypeError: print_error("Values must be strings, integers or floats.") # _ # _ _|_|___ _ _ _ ___ ___ # | | | | -_| | | | -_| _| # \_/|_|___|_____|___|_| def print_welcome(): header = """ ___ ___ __ __ | Y .---.-.----| |--|__.-----.-----. |. | _ | __| | | | -__| |. \\_/ |___._|____|__|__|__|__|__|_____|--. |: | | _ | | _ | _ | -__| _| |::.|:. |___._|__|__|___._|___ |_____|__| `--- ---' |_____| """ info = get_mm_info() print(bcolors.BOLD + bcolors.YELLOW + header + bcolors.ENDC) print(f"{bcolors.BOLD}{bcolors.BLUE}+ {bcolors.ENDC}{bcolors.BOLD}Pages: {bcolors.ENDC}{info['pages']}") print(f"{bcolors.BOLD}{bcolors.BLUE}+ {bcolors.ENDC}{bcolors.BOLD}Groups: {bcolors.ENDC}{len(info['groups'])}") print(f"{bcolors.BOLD}{bcolors.BLUE}+ {bcolors.ENDC}{bcolors.BOLD}Path: {bcolors.ENDC}{info['dirpath']}") print(f"{bcolors.BOLD}{bcolors.BLUE}+ {bcolors.ENDC}{bcolors.BOLD}Version: {bcolors.ENDC}{info['version']}") def print_table(data: dict): for k in KEY_ORDER: item = data[k] if k in data else "" print(f"{bcolors.BOLD}{bcolors.BLUE}{k}{bcolors.ENDC}\t{item}") def print_header(content: str): header = pyfiglet.figlet_format(content, font="cricket", width=TERM_SIZE.columns) print(bcolors.BOLD + bcolors.YELLOW + header + bcolors.ENDC) def print_page(name: str): try: with open(page_path(name), "r") as fp: page = json.load(fp) print_header(page["name"]) print_table(page) except FileNotFoundError: print_error(f"Page with name {name} does not exist.") def list_pages(): print_header("Pages") try: # List standalone Pages pages = get_standalone_pages() for page in pages: print(f"{bcolors.BOLD}{bcolors.BLUE}+ {bcolors.ENDC}{page[:-5]}") # List groups groups = get_groups() for group in groups: print(f"{bcolors.BOLD}{group}{bcolors.ENDC}") for dir_item, is_last in with_last(os.scandir(DIR_PATH + group)): prefix = "`--" if is_last else "|--" print(f"{bcolors.BOLD}{bcolors.BLUE}{prefix}+ {bcolors.ENDC}{dir_item.name[:-5]}") except FileNotFoundError: print_error("Page-Path corrupted, restart mm.") args = parse_args() init_function() match args["action"]: case "add": add_page(args=args) case "group": add_group(args["name"]) case "edit": edit_page(args=args) case "rm": delete_item(args["name"]) case "view": print_page(args["name"]) case "ls": list_pages() case "test": print(get_groups()) case _: print_welcome()