import os
import sys
import requests
from bs4 import BeautifulSoup
from rich.console import Console
from rich.table import Table
from colorama import Fore, init
from urllib.parse import urlparse
import base64
import dns.resolver

from argus.utils.util import clean_domain_input, clean_url, make_request, log_message
from argus.config.settings import API_KEYS

init(autoreset=True)
console = Console()

def banner():
    console.print(Fore.GREEN + """
    =============================================
          Argus - Advanced Malware & Phishing Check
    =============================================
    """)

def heuristic_analysis(url):
    suspicious_patterns = ["login", "secure", "account", "update", "verify", "signin", "admin"]
    console.print(Fore.YELLOW + "[*] Performing heuristic analysis...")
    parsed_url = urlparse(url)
    path = parsed_url.path.lower()
    if any(pattern in path for pattern in suspicious_patterns):
        console.print(Fore.RED + "[!] Suspicious keywords detected in URL path.")
    else:
        console.print(Fore.GREEN + "[+] No suspicious keywords found in URL path.")

def check_urlhaus(url):
    console.print(Fore.YELLOW + "[*] Checking URLhaus...")
    try:
        api_url = "https://urlhaus-api.abuse.ch/v1/url/"
        payload = {
            'url': url,
            'format': 'json'
        }
        headers = {
            'Content-Type': 'application/x-www-form-urlencoded'
        }
        response = requests.post(api_url, headers=headers, data=payload, timeout=10)
        if response and response.status_code == 200:
            data = response.json()
            if data.get("query_status") == "ok" and data.get("url_info"):
                console.print(Fore.RED + "[!] URL detected as malicious by URLhaus.")
                log_message("malware_phishing.log", "URL detected as malicious by URLhaus.")
            else:
                console.print(Fore.GREEN + "[+] URL not found in URLhaus database.")
                log_message("malware_phishing.log", "URL not found in URLhaus database.")
        else:
            console.print(Fore.RED + f"[!] URLhaus check failed with status code {response.status_code if response else 'No Response'}.")
    except Exception as e:
        console.print(Fore.RED + f"[!] Error during URLhaus check: {e}")

def check_blacklist_services(url):
    console.print(Fore.YELLOW + "[*] Checking against public blacklist services...")
    try:
        parsed_url = urlparse(url)
        domain = parsed_url.netloc
        ip_address = make_request(f"http://api.ipify.org?format=json")
        if ip_address:
            ip = ip_address.json().get("ip")
            if not ip:
                console.print(Fore.RED + "[!] Unable to retrieve IP address.")
                log_message("malware_phishing.log", "Unable to retrieve IP address.")
                return
            reversed_ip = '.'.join(ip.split('.')[::-1])
            dnsbl_domains = ["zen.spamhaus.org", "b.barracudacentral.org"]
            listed = False
            for dnsbl in dnsbl_domains:
                query = f"{reversed_ip}.{dnsbl}"
                try:
                    answers = dns.resolver.resolve(query, 'A')
                    for rdata in answers:
                        if rdata.address.startswith("127."):
                            console.print(Fore.RED + f"[!] IP {ip} is listed in {dnsbl}.")
                            log_message("malware_phishing.log", f"IP {ip} is listed in {dnsbl}.")
                            listed = True
                except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
                    continue
                except Exception as e:
                    console.print(Fore.RED + f"[!] DNS query error for {dnsbl}: {e}")
                    log_message("malware_phishing.log", f"DNS query error for {dnsbl}: {e}")
            if not listed:
                console.print(Fore.GREEN + "[+] IP not found in public blacklist services.")
                log_message("malware_phishing.log", "IP not found in public blacklist services.")
        else:
            console.print(Fore.RED + "[!] Unable to retrieve IP address for blacklist check.")
            log_message("malware_phishing.log", "Unable to retrieve IP address for blacklist check.")
    except Exception as e:
        console.print(Fore.RED + f"[!] Error during blacklist services check: {e}")
        log_message("malware_phishing.log", f"Error during blacklist services check: {e}")

def check_virustotal(url):
    console.print(Fore.YELLOW + "[*] Checking VirusTotal...")
    try:
        api_key = API_KEYS.get("VIRUSTOTAL_API_KEY")
        if not api_key:
            console.print(Fore.RED + "[!] VirusTotal API key not configured.")
            log_message("malware_phishing.log", "VirusTotal API key not configured.")
            return
        scan_url = "https://www.virustotal.com/api/v3/urls"
        encoded_url = base64.urlsafe_b64encode(url.encode()).decode().strip("=")
        payload = {
            'url': url
        }
        headers = {
            'x-apikey': api_key,
            'Content-Type': 'application/x-www-form-urlencoded'
        }
        response = requests.post(scan_url, headers=headers, data=payload, timeout=10)
        if response and response.status_code in [200, 201]:
            data = response.json()
            analysis_id = data['data']['id']
            analysis_url = f"https://www.virustotal.com/api/v3/analyses/{analysis_id}"
            analysis_response = requests.get(analysis_url, headers=headers, timeout=10)
            if analysis_response and analysis_response.status_code == 200:
                analysis_data = analysis_response.json()
                stats = analysis_data['data']['attributes']['stats']
                malicious = stats.get('malicious', 0)
                if malicious > 0:
                    console.print(Fore.RED + f"[!] VirusTotal detected {malicious} malicious sources.")
                    log_message("malware_phishing.log", f"VirusTotal detected {malicious} malicious sources.")
                else:
                    console.print(Fore.GREEN + "[+] VirusTotal found no malicious sources.")
                    log_message("malware_phishing.log", "VirusTotal found no malicious sources.")
            else:
                console.print(Fore.RED + f"[!] VirusTotal analysis failed with status code {analysis_response.status_code if analysis_response else 'No Response'}.")
                log_message("malware_phishing.log", f"VirusTotal analysis failed with status code {analysis_response.status_code if analysis_response else 'No Response'}.")
        else:
            console.print(Fore.RED + f"[!] VirusTotal scan failed with status code {response.status_code if response else 'No Response'}.")
            log_message("malware_phishing.log", f"VirusTotal scan failed with status code {response.status_code if response else 'No Response'}.")
    except Exception as e:
        console.print(Fore.RED + f"[!] Error during VirusTotal check: {e}")
        log_message("malware_phishing.log", f"Error during VirusTotal check: {e}")

def additional_heuristic_checks(url):
    console.print(Fore.YELLOW + "[*] Performing additional heuristic checks...")
    try:
        parsed_url = urlparse(url)
        if parsed_url.scheme != "https":
            console.print(Fore.RED + "[!] URL does not use HTTPS. Potential security risk.")
            log_message("malware_phishing.log", "URL does not use HTTPS.")
        else:
            console.print(Fore.GREEN + "[+] URL uses HTTPS.")
            log_message("malware_phishing.log", "URL uses HTTPS.")
    except Exception as e:
        console.print(Fore.RED + f"[!] Error during additional heuristic checks: {e}")
        log_message("malware_phishing.log", f"Error during additional heuristic checks: {e}")

def check_malware_phishing(url):
    heuristic_analysis(url)
    additional_heuristic_checks(url)
    check_urlhaus(url)
    check_blacklist_services(url)
    check_virustotal(url)

def display_malware_phishing():
    console.print(Fore.CYAN + "\n[*] Malware & Phishing checks completed.\n")

def main(target):
    banner()
    target = clean_url(clean_domain_input(target))
    console.print(Fore.WHITE + f"[*] Checking for malware and phishing threats for: {target}")
    check_malware_phishing(target)
    display_malware_phishing()

if __name__ == "__main__":
    if len(sys.argv) > 1:
        try:
            target = sys.argv[1]
            main(target)
        except KeyboardInterrupt:
            console.print(Fore.RED + "\n[!] Process interrupted by user.")
            sys.exit(1)
    else:
        console.print(Fore.RED + "[!] No target provided. Please pass a domain or URL.")
        sys.exit(1)
