import os
from colorama import init, Fore
from malware_sentinel.api_config import get_api_key
from malware_sentinel.utils import (
    display_title, list_and_get_the_file, if_is_valid_md5,
    request_page, print_analysis_results,
    is_valid_ip, is_valid_hash,is_valid_domain,check_OS_Execute
)
init(autoreset=True)

def clear():
    if os.name == 'nt':
        os.system('cls')
    else:
        os.system('clear')
def main():
    try:
        display_title()

        api_key = get_api_key()
        if not api_key:
            print(Fore.RED + "Please provide a valid API key.")
            return

        while True:
            print(Fore.CYAN + "\nOptions: 'file', 'hash', 'ip', 'domain', 'exit'")
            choice = input(Fore.CYAN + "What do you want to scan? ").strip().lower()

            if choice == "exit":
                print(Fore.GREEN + "Exiting program. Goodbye!")
                return   #  exit program 

            elif choice == 'file':
                filePrompt = input(Fore.CYAN + "Enter file path to scan (or 'exit' to quit): ").strip()
                if filePrompt.lower() == "exit":
                    print(Fore.GREEN + "Exiting program. Goodbye!")
                    return

                file_path = list_and_get_the_file(filePrompt)
                if not file_path:
                    print(Fore.RED + "File not found.")
                    continue

                md5sum = if_is_valid_md5(file_path)
                if not md5sum:
                    print(Fore.RED + "File could not be read or is invalid.")
                    continue

                print(Fore.CYAN + f"MD5 checksum: {md5sum}")
                url = f'https://www.virustotal.com/api/v3/files/{md5sum}'
                data = request_page(url, api_key, "file")
                if data:
                    print_analysis_results(
                        data['data']['attributes']['last_analysis_stats'],
                        data['data']['attributes']['total_votes'],
                        md5sum,
                        'file',
                        data
                    )
                else:
                    check_OS_Execute(filePrompt)
            elif choice == 'hash':
                hash_input = input(Fore.CYAN + "Enter hash to scan (or 'exit' to quit): ").strip()
                if hash_input.lower() == "exit":
                    print(Fore.GREEN + "Exiting program. Goodbye!")
                    return

                if not is_valid_hash(hash_input):
                    print(Fore.RED + "Invalid hash format. Must be MD5, SHA1, or SHA256.")
                    continue

                url = f'https://www.virustotal.com/api/v3/files/{hash_input}'
                data = request_page(url, api_key, "hash")
                if data:
                    print_analysis_results(
                        data['data']['attributes']['last_analysis_stats'],
                        data['data']['attributes']['total_votes'],
                        hash_input,
                        'hash',
                        data
                    )
                else:
                    check_OS_Execute(filePrompt)

            elif choice == 'ip':
                ip = input(Fore.CYAN + "Enter IP address to scan (or 'exit' to quit): ").strip()
                if ip.lower() == "exit":
                    print(Fore.GREEN + "Exiting program. Goodbye!")
                    return

                if not is_valid_ip(ip):
                    print(Fore.RED + "Invalid IP format. Please enter a valid IPv4 or IPv6 address.")
                    continue

                url = f'https://www.virustotal.com/api/v3/ip_addresses/{ip}'
                data = request_page(url, api_key, "ip")
                if data:
                    print_analysis_results(
                        data['data']['attributes']['last_analysis_stats'],
                        data['data']['attributes']['total_votes'],
                        ip,
                        'ip',
                        data
                    )

            elif choice == 'domain':
                domain = input(Fore.CYAN + "Enter domain to scan (or 'exit' to quit): ").strip()
                if domain.lower() == "exit":
                    print(Fore.GREEN + "Exiting program. Goodbye!")
                    return

                if not is_valid_domain(domain):
                    print(Fore.RED + "Invalid domain format. Please enter something like 'example.com'.")
                    continue

                url = f'https://www.virustotal.com/api/v3/domains/{domain}'
                data = request_page(url, api_key, "domain")
                if data:
                    print_analysis_results(
                        data['data']['attributes']['last_analysis_stats'],
                        data['data']['attributes']['total_votes'],
                        domain,
                        'domain',
                        data
                    )

            else:
                print(Fore.RED + "Invalid choice. Please enter 'file', 'hash', 'ip', 'domain', or 'exit'.")
    except KeyboardInterrupt:
        clear()
if __name__ == "__main__":
    main()
