VT_IOC_Analyzer.py

import os
import csv
import time
import requests
import pandas as pd
import numpy as np
from collections import defaultdict
import urllib.parse
import base64
from concurrent.futures import ThreadPoolExecutor, as_completed
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.cluster import KMeans
import json
import logging
from requests.adapters import HTTPAdapter
from requests.packages.urllib3.util.retry import Retry
import ipaddress

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Global variables
last_request_time = 0
BATCH_SIZE = 10
MAX_RETRIES = 3
BACKOFF_FACTOR = 0.3

def load_iocs(file_path, remove_duplicates=False):
    iocs = defaultdict(set)  # Changed to set to automatically remove duplicates
    with open(file_path, 'r') as f:
        for line in f:
            parts = line.strip().split('\t')
            if len(parts) >= 2:
                ioc_type = parts[0]
                ioc_value = parts[1]
                if ioc_type.startswith('FileHash-'):
                    hash_type = ioc_type.split('-')[1].lower()
                    iocs[hash_type].add(ioc_value)
                elif ioc_type == 'IPv4':
                    iocs['ip_addresses'].add(ioc_value)
                elif ioc_type == 'URL':
                    iocs['urls'].add(ioc_value)
                elif ioc_type == 'domain':
                    iocs['domain_names'].add(ioc_value)
    
    # Convert sets back to lists
    return {k: list(v) for k, v in iocs.items()}

def create_session_with_retry():
    session = requests.Session()
    retries = Retry(total=MAX_RETRIES,
                    backoff_factor=BACKOFF_FACTOR,
                    status_forcelist=[429, 500, 502, 503, 504])
    session.mount('https://', HTTPAdapter(max_retries=retries))
    return session

def rate_limited_request(session, url, headers):
    global last_request_time
    current_time = time.time()
    if current_time - last_request_time < 15:
        time.sleep(15 - (current_time - last_request_time))
    
    response = session.get(url, headers=headers)
    last_request_time = time.time()
    return response

def get_yara_rule_names(data):
    yara_rules = data['data']['attributes'].get('crowdsourced_yara_results', [])
    return ', '.join(rule['rule_name'] for rule in yara_rules if rule.get('rule_name'))

def analyze_ioc_batch(iocs, ioc_type, api_key, consolidate_file_hashes=False, processed_files=None):
    if processed_files is None:
        processed_files = set()

    base_url = "https://www.virustotal.com/api/v3/"
    headers = {"x-apikey": api_key}
    session = create_session_with_retry()
    results = []

    for ioc in iocs:
        if ioc_type in ['md5', 'sha1', 'sha256', 'sha512']:
            if consolidate_file_hashes and ioc in processed_files:
                continue
            url = f"{base_url}files/{ioc}"
        elif ioc_type == 'ip_addresses':
            url = f"{base_url}ip_addresses/{ioc}"
        elif ioc_type == 'urls':
            url_id = urllib.parse.quote(base64.urlsafe_b64encode(ioc.encode()).decode().strip("="))
            url = f"{base_url}urls/{url_id}"
        else:  # domain_names
            url = f"{base_url}domains/{ioc}"
        
        try:
            response = rate_limited_request(session, url, headers)
            response.raise_for_status()
            data = response.json()
            malicious_count = data['data']['attributes']['last_analysis_stats']['malicious']
            first_seen = data['data']['attributes'].get('first_submission_date')
            last_seen = data['data']['attributes'].get('last_analysis_date')
            threat_name = get_yara_rule_names(data)
            file_name = data['data']['attributes'].get('meaningful_name', 'N/A')
            file_type = data['data']['attributes'].get('type_description', 'N/A')
            country = data['data']['attributes'].get('country', 'Unknown')
            asn = data['data']['attributes'].get('as_owner', 'Unknown')

            if ioc_type in ['md5', 'sha1', 'sha256', 'sha512'] and consolidate_file_hashes:
                file_hashes = {
                    data['data']['attributes'].get('md5'),
                    data['data']['attributes'].get('sha1'),
                    data['data']['attributes'].get('sha256'),
                    data['data']['attributes'].get('sha512')
                }
                processed_files.update(file_hashes)

            results.append((ioc, ioc_type, malicious_count, first_seen, last_seen, threat_name, file_name, file_type, country, asn, data))
        except requests.exceptions.RequestException as e:
            logging.error(f"Error analyzing {ioc}: {str(e)}")
            results.append((ioc, ioc_type, 0, None, None, 'N/A', 'N/A', 'N/A', 'Unknown', 'Unknown', {}))

    return results, processed_files

def get_color(detection_count):
    if detection_count == 0:
        return 'green'
    elif 1 <= detection_count <= 5:
        return 'yellow'
    elif 6 <= detection_count <= 10:
        return 'orange'
    else:
        return 'red'

def generate_csv_report(results, filename='ioc_report.csv'):
    with open(filename, 'w', newline='') as csvfile:
        fieldnames = ['IOC', 'Type', 'Detection Count', 'Color', 'Country', 'ASN', 'First Seen', 'Last Seen', 'Threat Name', 'File Name', 'File Type']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
        for result in results:
            writer.writerow({
                'IOC': result[0],
                'Type': result[1],
                'Detection Count': result[2],
                'Color': get_color(result[2]),
                'Country': result[8],
                'ASN': result[9],
                'First Seen': time.strftime('%Y-%m-%d %H:%M:%S', time.gmtime(result[3])) if result[3] else 'N/A',
                'Last Seen': time.strftime('%Y-%m-%d %H:%M:%S', time.gmtime(result[4])) if result[4] else 'N/A',
                'Threat Name': result[5] if result[5] else 'N/A',
                'File Name': result[6],
                'File Type': result[7]
            })
        
        # Write intelligence summary
        csvfile.write("\n")  # Add a new line before the summary
        csvfile.write("Intelligence Summary\n")
        csvfile.write("Total IOCs, Malicious IOCs, Average Detection Count, Max Detection Count, Top Threat Names\n")
        
        total_iocs = len(results)
        malicious_iocs = sum(1 for r in results if r[2] > 0)
        average_detection_count = sum(r[2] for r in results) / total_iocs if total_iocs > 0 else 0
        max_detection_count = max(r[2] for r in results) if results else 0
        
        # Count the top threat names
        threat_names = defaultdict(int)
        for result in results:
            if result[5] and result[5] != 'N/A':
                for threat in result[5].split(', '):
                    threat_names[threat] += 1
        
        top_threat_names = sorted(threat_names.items(), key=lambda x: x[1], reverse=True)[:5]
        top_threat_names_str = ', '.join(f"{name} ({count})" for name, count in top_threat_names)
        
        csvfile.write(f"{total_iocs}, {malicious_iocs}, {average_detection_count:.2f}, {max_detection_count}, {top_threat_names_str}\n")
    
    logging.info(f"CSV report generated: {filename}")

def perform_statistical_analysis(results):
    df = pd.DataFrame([r[:10] for r in results], columns=['IOC', 'Type', 'Detection Count', 'First Seen', 'Last Seen', 'Threat Name', 'File Name', 'File Type', 'Country', 'ASN'])
    stats = {
        'Total IOCs': len(df),
        'Malicious IOCs': len(df[df['Detection Count'] > 0]),
        'Average Detection Count': df['Detection Count'].mean(),
        'Max Detection Count': df['Detection Count'].max(),
        'IOCs by Type': df['Type'].value_counts().to_dict(),
        'Top 5 IOCs by Detection Count': df.nlargest(5, 'Detection Count')[['IOC', 'Type', 'Detection Count']].to_dict('records'),
        'Top 5 Countries': df['Country'].value_counts().nlargest(5).to_dict(),
        'Top 5 ASNs': df['ASN'].value_counts().nlargest(5).to_dict(),
        'Top 5 Threat Names': df['Threat Name'].apply(lambda x: x.split(', ') if x != 'N/A' else []).explode().value_counts().nlargest(5).to_dict(),
        'Top 5 File Types': df['File Type'].value_counts().nlargest(5).to_dict()
    }
    return stats

def perform_text_analysis(results):
    text_data = [list(ioc) for ioc in results if ioc[1] in ['domain_names', 'urls']]
    
    if not text_data:
        return "No domain names or URLs to analyze.", {}

    vectorizer = CountVectorizer()
    X = vectorizer.fit_transform([item[0] for item in text_data])

    n_clusters = min(5, len(text_data))
    kmeans = KMeans(n_clusters=n_clusters)
    kmeans.fit(X)

    for i, item in enumerate(text_data):
        item.append(int(kmeans.labels_[i]))

    cluster_terms = {}
    feature_names = vectorizer.get_feature_names_out()
    for i in range(n_clusters):
        center = kmeans.cluster_centers_[i]
        top_terms = [feature_names[j] for j in center.argsort()[-5:][::-1]]
        cluster_terms[int(i)] = top_terms

    return text_data, cluster_terms

def process_iocs(iocs, api_key, consolidate_file_hashes=False):
    results = []
    processed_files = set()
    with ThreadPoolExecutor(max_workers=5) as executor:
        future_to_batch = {}
        for ioc_type, ioc_list in iocs.items():
            for i in range(0, len(ioc_list), BATCH_SIZE):
                batch = ioc_list[i:i+BATCH_SIZE]
                future = executor.submit(analyze_ioc_batch, batch, ioc_type, api_key, consolidate_file_hashes, processed_files)
                future_to_batch[future] = (batch, ioc_type)
        
        for future in as_completed(future_to_batch):
            batch, ioc_type = future_to_batch[future]
            try:
                batch_results, processed_files = future.result()
                results.extend(batch_results)
            except Exception as e:
                logging.error(f"Error processing batch: {str(e)}")
    
    return results

def get_file_path():
    while True:
        file_path = input("Please enter the path to your IOC file: ").strip()
        if os.path.exists(file_path):
            return file_path
        else:
            logging.warning("The specified file does not exist. Please try again.")

def get_api_key():
    return input("Please enter your VirusTotal API key: ").strip()

def json_serial(obj):
    if isinstance(obj, (np.integer, np.floating, np.bool_)):
        return obj.item()
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    if isinstance(obj, pd.DataFrame):
        return obj.to_dict(orient='records')
    if isinstance(obj, pd.Series):
        return obj.to_dict()
    raise TypeError(f"Type {type(obj)} not serializable")

def convert_to_serializable(data):
    if isinstance(data, dict):
        return {k: convert_to_serializable(v) for k, v in data.items()}
    elif isinstance(data, list):
        return [convert_to_serializable(v) for v in data]
    elif isinstance(data, (np.integer, np.floating, np.bool_)):
        return data.item()
    elif isinstance(data, np.ndarray):
        return data.tolist()
    elif isinstance(data, pd.DataFrame):
        return data.to_dict(orient='records')
    elif isinstance(data, pd.Series):
        return data.to_dict()
    else:
        return data

def save_json_report(data, filename='ioc_analysis_report.json'):
    serializable_data = convert_to_serializable(data)
    with open(filename, 'w') as f:
        json.dump(serializable_data, f, indent=4, default=json_serial)
    logging.info(f"JSON report generated: {filename}")

def analyze_ip_ranges(results):
    ip_ranges = defaultdict(list)
    for result in results:
        if result[1] == 'ip_addresses':
            ip = ipaddress.ip_address(result[0])
            for cidr in ['8', '16', '24']:
                network = ipaddress.ip_network(f"{ip}/{cidr}", strict=False)
                ip_ranges[str(network)].append(result)
    
    return {range: len(ips) for range, ips in ip_ranges.items() if len(ips) > 1}

# Main execution
if __name__ == "__main__":
    file_path = get_file_path()
    api_key = get_api_key()

    remove_duplicates = input("Do you want to remove duplicates from the input? (y/n): ").lower() == 'y'
    consolidate_file_hashes = input("Do you want to consolidate results for different checksums of the same file? (y/n): ").lower() == 'y'

    iocs = load_iocs(file_path, remove_duplicates)
    logging.info("Processing IOCs... This may take a while depending on the number of IOCs.")
    results = process_iocs(iocs, api_key, consolidate_file_hashes)

    # Generate CSV report
    generate_csv_report(results)

    # Perform statistical analysis
    stats = perform_statistical_analysis(results)
    logging.info("\nStatistical Analysis:")
    for key, value in stats.items():
        logging.info(f"{key}: {value}")

    # Perform text-based analysis
    text_analysis, cluster_terms = perform_text_analysis(results)
    logging.info("\nText-Based Analysis:")
    logging.info(f"Clustered {len(text_analysis)} domain names and URLs into groups.")
    for cluster_id, terms in cluster_terms.items():
        logging.info(f"\nCluster {cluster_id} (Top terms: {', '.join(terms)}):")
        for item in text_analysis:
            if item[-1] == cluster_id:
                logging.info(f"  {item[0]} ({item[1]}) - Detections: {item[2]}, Color: {get_color(item[2])}")

    # Analyze IP ranges
    ip_range_analysis = analyze_ip_ranges(results)
    logging.info("\nIP Range Analysis:")
    for range, count in ip_range_analysis.items():
        logging.info(f"Range {range}: {count} IPs")

    # Save comprehensive report as JSON
    full_report = {
        'statistics': stats,
        'text_analysis': {
            'clustered_data': text_analysis,
            'cluster_terms': cluster_terms
        },
        'ip_range_analysis': ip_range_analysis,
        'raw_results': results
    }
    save_json_report(full_report)

    logging.info("\nAnalysis complete. Check the CSV and JSON reports for detailed results.")