mm/mm.py

import argparse
import json
import os
import pyfiglet 
import sys
import textwrap

#    _     _ _   
#   |_|___|_| |_ 
#   | |   | |  _|
#   |_|_|_|_|_|  

SCRIPT_VERSION = "0.1"

class bcolors:
    GREY = '\033[90m'      
    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", "desc", "ipv4", "ipv6", "os", "ram", "cores"] 
# 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, group: str=None) -> str:
    if group and os.path.exists(DIR_PATH + group):
        return DIR_PATH + group +  f"/{hostname}.json"
    else:
        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("-4","--ipv4", required=False)
    parser.add_argument("-6","--ipv6", required=False)
    parser.add_argument("-o","--os", required=False)
    parser.add_argument("-r","--ram", required=False)
    parser.add_argument("-c","--cores", required=False)
    parser.add_argument("-d","--desc", required=False)
    parser.add_argument("-g","--group", 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 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"], args["group"]), "r") as fp:
            page = json.load(fp)
        page.update(data)
        with open(page_path(data["name"], args["group"]), "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"], args["group"])) and not args["force"]:
            print_error(f"Page {data['name']} already exists, use --force to overwrite.")
            exit()
        with open(page_path(data["name"], args["group"]), "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, group: str):
    try: 
        with open(page_path(name, group), "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 sorted(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 in sorted(os.listdir(DIR_PATH + group)):
                with open(page_path(dir_item[:-5], group), "r") as fp:
                    content = json.load(fp)
                desc = content["desc"] if "desc" in content else ""
                ram = content["ram"] if "ram" in content else ""
                print(f"{bcolors.BOLD}{bcolors.BLUE}+ {bcolors.ENDC}{dir_item[:-5]}\t{bcolors.RED}{desc}\t{bcolors.GREY}{ram}{bcolors.ENDC}")
    except FileNotFoundError as e:
        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"], args["group"])
    case "ls": list_pages()
    case "test": print(args)
    case _: print_welcome()