#!/usr/bin/env bash

# Check if a filename was provided
if [ "$#" -ne 1 ]; then
    echo "Usage: $0 <nmap_xml_file>"
    exit 1
fi

XML_FILE="$1"
DB_NAME="pentest_results.db"

# Check if the file exists
if [ ! -f "$XML_FILE" ]; then
    echo "Error: File $XML_FILE does not exist."
    exit 1
fi

# Function to sanitize input for SQL
sanitize() {
    echo "${1//\'/\'\'}"
}

# Process the XML file
xmlstarlet sel -T -t -m "//host" \
    -v "address[@addrtype='ipv4']/@addr" -o "|" \
    -v "hostnames/hostname/@name" -o "|" \
    -m "ports/port" \
        -v "concat(@portid, ',')" \
        -v "concat(@protocol, ',')" \
        -v "concat(service/@name, ',')" \
        -v "concat(service/@version, ',')" \
        -b -o "|" \
    -n \
"$XML_FILE" | while IFS='|' read -r ip hostname ports_info; do
    # Skip empty lines
    [ -z "$ip" ] && continue

    # Process each port for the current IP
    IFS=',' read -ra port_array <<< "$ports_info"
    for ((i=0; i<${#port_array[@]}; i+=4)); do
        port="${port_array[i]}"
        protocol="${port_array[i+1]}"
        service="${port_array[i+2]}"
        version="${port_array[i+3]}"

        # Sanitize inputs
        ip=$(sanitize "$ip")
        hostname=$(sanitize "$hostname")
        port=$(sanitize "$port")
        protocol=$(sanitize "$protocol")
        service=$(sanitize "$service")
        version=$(sanitize "$version")

        # Insert or replace the data
        sqlite3 "$DB_NAME" <<EOF
INSERT OR REPLACE INTO nmap_scans (ip_address, hostname, port, protocol, service, version)
VALUES ('$ip', '$hostname', '$port', '$protocol', '$service', '$version');
EOF
    done
done

# Add indexes for better query performance
sqlite3 "$DB_NAME" <<EOF
CREATE INDEX IF NOT EXISTS idx_ip_address ON nmap_scans (ip_address);
CREATE INDEX IF NOT EXISTS idx_port ON nmap_scans (port);
EOF

echo "Data import completed."
