import pandas as pd
import os
import whois
import socket
import requests
from ipwhois import IPWhois

def aggregate_csv_to_excel(folder_path, output_excel):
    with pd.ExcelWriter(output_excel, engine='xlsxwriter') as writer:
        for filename in os.listdir(folder_path):
            if filename.endswith('.csv'):
                file_path = os.path.join(folder_path, filename)
                df = pd.read_csv(file_path)
                
                # Filter relevant columns
                columns = ['timestamp', 'client_ip', 'server_ip', 'server_port', 'direction', 'total_bytes', 'country', 'flow_state']
                if 'dns' in df.columns:
                    columns.append('dns')
                df = df[columns]
                
                # Investigate each server IP
                df['whois_info'] = df['server_ip'].apply(get_whois_info)
                df['geo_info'] = df['server_ip'].apply(get_geo_info)
                df['virus_total'] = df['server_ip'].apply(check_virus_total)
                df['dns_lookup'] = df['server_ip'].apply(dns_lookup)

                # Write to Excel
                sheet_name = os.path.splitext(filename)[0]
                df.to_excel(writer, sheet_name=sheet_name, index=False)

def get_whois_info(ip):
    try:
        w = whois.whois(ip)
        return f"Registrar: {w.registrar}, Country: {w.country}"
    except Exception as e:
        return f"WHOIS Error: {str(e)}"

def get_geo_info(ip):
    try:
        obj = IPWhois(ip)
        res = obj.lookup_rdap()
        return f"Country: {res['asn_country_code']}, ASN: {res['asn']}"
    except Exception as e:
        return f"Geo Error: {str(e)}"

def check_virus_total(ip):
    try:
        url = f"https://www.virustotal.com/api/v3/ip_addresses/{ip}"
        headers = {"x-apikey": "your-api-key-here"}  # Replace with your VirusTotal API key
        response = requests.get(url, headers=headers)
        if response.status_code == 200:
            data = response.json()
            return f"Reputation: {data.get('data', {}).get('attributes', {}).get('reputation', 'N/A')}"
        else:
            return f"VT Error: {response.status_code}"
    except Exception as e:
        return f"VT Error: {str(e)}"

def dns_lookup(ip):
    try:
        return socket.gethostbyaddr(ip)[0]
    except socket.herror:
        return "DNS Lookup Error"

def main():
    current_dir = os.getcwd()
    user_input = input(f"Do you want to use the current directory ({current_dir}) to pull CSV files? (yes/no): ").strip().lower()
    
    if user_input == 'yes':
        folder_path = current_dir
    else:
        folder_path = input("Please enter the directory where the CSV files are located: ").strip()
    
    output_excel = os.path.join(current_dir, 'aggregated_investigation.xlsx')
    aggregate_csv_to_excel(folder_path, output_excel)
    print(f"Aggregated investigation saved to {output_excel}")

if __name__ == "__main__":
    main()