#!/usr/bin/env bash

DB_NAME="pentest_results.db"

# Function to query the database and format results
query_db() {
    sqlite3 -header -csv $DB_NAME "$1"
}

# Function to escape special characters for LaTeX
escape_latex() {
    sed -e 's/\\/\\textbackslash/g' \
        -e 's/\^/\\textasciicircum/g' \
        -e 's/_/\\_/g' \
        -e 's/~/\\textasciitilde/g' \
        -e 's/&/\\&/g' \
        -e 's/%/\\%/g' \
        -e 's/#/\\#/g' \
        -e 's/\$/\\$/g' \
        -e 's/{/\\{/g' \
        -e 's/}/\\}/g'
}

# Get all unique IP addresses with vulnerabilities
ip_addresses=$(query_db "SELECT DISTINCT ip_address FROM nmap_scans WHERE vulnerability IS NOT NULL AND vulnerability != '';" | tail -n +2)

# Create LaTeX content
create_latex_content() {
    echo "\\documentclass{article}"
    echo "\\usepackage[margin=1in]{geometry}"
    echo "\\usepackage{longtable}"
    echo "\\usepackage{array}"
    echo "\\usepackage{booktabs}"
    echo "\\usepackage{multirow}"
    echo "\\usepackage{xcolor}"
    echo "\\begin{document}"
    echo "\\title{Penetration Test Report}"
    echo "\\author{Your Name}"
    echo "\\date{\\today}"
    echo "\\maketitle"
    
    echo "\\section{Scan Results}"
    
    IFS=$'\n'
    for ip in $ip_addresses; do
        escaped_ip=$(echo "$ip" | escape_latex)
        echo "\\subsection{IP Address: $escaped_ip}"
        echo "\\begin{longtable}{p{0.15\\textwidth}p{0.15\\textwidth}p{0.15\\textwidth}p{0.25\\textwidth}p{0.3\\textwidth}}"
        echo "\\toprule"
        echo "\\textbf{Port} & \\textbf{Protocol} & \\textbf{Service} & \\textbf{Version} & \\textbf{Vulnerability} \\\\"
        echo "\\midrule"
        echo "\\endhead"
        
        query_db "SELECT hostname, ip_address, port, protocol, service, version, vulnerability 
                  FROM nmap_scans 
                  WHERE ip_address='$ip' AND vulnerability IS NOT NULL AND vulnerability != '';" | tail -n +2 | while IFS=',' read -r hostname ip port protocol service version vulnerability; do
            escaped_hostname=$(echo "$hostname" | escape_latex)
            escaped_service=$(echo "$service" | escape_latex)
            escaped_version=$(echo "$version" | escape_latex)
            escaped_vulnerability=$(echo "$vulnerability" | escape_latex)
            
            if [ -z "$escaped_hostname" ]; then
                escaped_hostname="N/A"
            fi
            
            echo "\\multicolumn{5}{l}{\\textbf{Hostname:} $escaped_hostname} \\\\"
            echo "$port & $protocol & $escaped_service & $escaped_version & $escaped_vulnerability \\\\"
            echo "\\midrule"
        done
        
        echo "\\bottomrule"
        echo "\\end{longtable}"
    done
    
    echo "\\end{document}"
}

# Generate LaTeX file
create_latex_content > pentest_report.tex

# Compile LaTeX to PDF
pdflatex -interaction=nonstopmode pentest_report.tex
