Hunt-AI/Modules/Submenus/analyst_notebook.py

import json
from datetime import datetime

# In-memory data store
data_store = {
    "ips": [],
    "domains": [],
    "services": [],
    "notes": [],
}

# File to store the data persistently
DATA_FILE = "analyst_notebook.json"

def save_data():
    """
    Saves the current data store to the JSON file.
    """
    try:
        with open(DATA_FILE, "w") as file:
            json.dump(data_store, file, indent=4)
        print(f"Data successfully saved to {DATA_FILE}.")
    except Exception as e:
        print(f"Error saving data: {e}")

def load_data():
    """
    Loads the data store from the JSON file.
    """
    global data_store
    try:
        with open(DATA_FILE, "r") as file:
            data_store = json.load(file)
        print(f"Data successfully loaded from {DATA_FILE}.")
    except FileNotFoundError:
        print(f"{DATA_FILE} not found. Starting with a new data store.")
    except Exception as e:
        print(f"Error loading data: {e}")

def add_entry(entry_type, value):
    """
    Adds a new entry with a timestamp to the specified entry type.
    """
    timestamped_value = {"value": value, "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
    data_store[entry_type].append(timestamped_value)
    save_data()  # Save data after modification
    print(f"{entry_type.capitalize()} added: {value} (at {timestamped_value['timestamp']})")

def view_entries(entry_type):
    """
    Displays all entries of the specified type, including timestamps.
    """
    entries = data_store[entry_type]
    if entries:
        print(f"Listing all {entry_type.capitalize()}:")
        for idx, entry in enumerate(entries, start=1):
            print(f"{idx}. {entry['value']} (Added on: {entry['timestamp']})")
    else:
        print(f"No {entry_type} found.")

def modify_entry(entry_type, index, new_value):
    """
    Modifies an existing entry and updates its timestamp.
    """
    try:
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        data_store[entry_type][index] = {"value": new_value, "timestamp": timestamp}
        save_data()  # Save data after modification
        print(f"{entry_type.capitalize()} updated to: {new_value} (at {timestamp})")
    except IndexError:
        print(f"Invalid index for {entry_type}.")

def delete_entry(entry_type, index):
    """
    Deletes an entry by index from the specified entry type.
    """
    try:
        removed = data_store[entry_type].pop(index)
        save_data()  # Save data after modification
        print(f"Removed {entry_type.capitalize()}: {removed['value']}")
    except IndexError:
        print(f"Invalid index for {entry_type}.")

def show_all_notes():
    """
    Displays all entries from all categories (IPs, Domains, Services, Notes).
    """
    print("\nAll Entries (with Timestamps):")
    for category, entries in data_store.items():
        if entries:
            print(f"\n{category.capitalize()}:")
            for idx, entry in enumerate(entries, start=1):
                print(f"{idx}. {entry['value']} (Added on: {entry['timestamp']})")
        else:
            print(f"\nNo {category.capitalize()} found.")

def analyst_notebook_submenu():
    """
    Main menu for the analyst notebook.
    """
    while True:
        print("\nAnalyst Notebook")
        print("[N] Notes")
        print("[IP] IPs")
        print("[D] Domains")
        print("[S] Services")
        print("[A] Show All Entries")
        print("[0] Return to Main Menu")

        choice = input("Enter your choice: ").strip().lower()

        if choice == "ip":
            manage_type("ips")
        elif choice == "d":
            manage_type("domains")
        elif choice == "s":
            manage_type("services")
        elif choice == "n":
            manage_type("notes")
        elif choice == "a":
            show_all_notes()
        elif choice == "0":
            break
        else:
            print("Invalid choice.")

def manage_type(entry_type):
    """
    Handles management of a specific type (IP, domain, service, or note).
    """
    while True:
        print(f"\nManaging {entry_type.capitalize()}")
        print("[1] Add")
        print("[2] View")
        print("[3] Modify")
        print("[4] Delete")
        print("[0] Back to Analyst Notebook")

        choice = input("Enter your choice: ").strip()

        if choice == "1":
            value = input(f"Enter {entry_type[:-1]} to add: ").strip()
            add_entry(entry_type, value)
        elif choice == "2":
            view_entries(entry_type)
        elif choice == "3":
            index = int(input(f"Enter the index to modify: ")) - 1
            new_value = input(f"Enter new {entry_type[:-1]}: ").strip()
            modify_entry(entry_type, index, new_value)
        elif choice == "4":
            index = int(input(f"Enter the index to delete: ")) - 1
            delete_entry(entry_type, index)
        elif choice == "0":
            break
        else:
            print("Invalid choice.")