dnsrecon/src/main.py

# File: src/main.py
"""Main CLI interface for the reconnaissance tool."""

import click
import json
import sys
import logging
from pathlib import Path
from .config import Config
from .reconnaissance import ReconnaissanceEngine
from .report_generator import ReportGenerator
from .web_app import create_app

# Module logger
logger = logging.getLogger(__name__)

@click.command()
@click.argument('target', required=False)
@click.option('--web', is_flag=True, help='Start web interface instead of CLI')
@click.option('--shodan-key', help='Shodan API key')
@click.option('--virustotal-key', help='VirusTotal API key')
@click.option('--max-depth', default=2, help='Maximum recursion depth (default: 2)')
@click.option('--output', '-o', help='Output file prefix (will create .json and .txt files)')
@click.option('--json-only', is_flag=True, help='Only output JSON')
@click.option('--text-only', is_flag=True, help='Only output text report')
@click.option('--port', default=5000, help='Port for web interface (default: 5000)')
@click.option('--verbose', '-v', is_flag=True, help='Enable verbose logging (DEBUG level)')
@click.option('--quiet', '-q', is_flag=True, help='Quiet mode (WARNING level only)')
def main(target, web, shodan_key, virustotal_key, max_depth, output, json_only, text_only, port, verbose, quiet):
    """DNS Reconnaissance Tool
    
    Examples:
        recon example.com                    # Scan example.com
        recon example                        # Try example.* for all TLDs
        recon example.com --max-depth 3      # Deeper recursion
        recon example.com -v                 # Verbose logging
        recon --web                          # Start web interface
    """
    
    # Determine log level
    if verbose:
        log_level = "DEBUG"
    elif quiet:
        log_level = "WARNING"
    else:
        log_level = "INFO"
    
    # Create configuration and setup logging
    config = Config.from_args(shodan_key, virustotal_key, max_depth, log_level)
    config.setup_logging(cli_mode=True)
    
    if web:
        # Start web interface
        logger.info("🌐 Starting web interface...")
        app = create_app(config)
        logger.info(f"🚀 Web interface starting on http://0.0.0.0:{port}")
        app.run(host='0.0.0.0', port=port, debug=False)  # Changed debug to False to reduce noise
        return
    
    if not target:
        click.echo("❌ Error: TARGET is required for CLI mode. Use --web for web interface.")
        sys.exit(1)
    
    # Initialize reconnaissance engine
    logger.info("🔧 Initializing reconnaissance engine...")
    engine = ReconnaissanceEngine(config)
    
    # Set up progress callback for CLI
    def progress_callback(message, percentage=None):
        if percentage is not None:
            click.echo(f"[{percentage:3d}%] {message}")
        else:
            click.echo(f"       {message}")
    
    engine.set_progress_callback(progress_callback)
    
    # Display startup information
    click.echo("=" * 60)
    click.echo("🔍 DNS RECONNAISSANCE TOOL")
    click.echo("=" * 60)
    click.echo(f"🎯 Target: {target}")
    click.echo(f"📊 Max recursion depth: {max_depth}")
    click.echo(f"🌐 DNS servers: {', '.join(config.DNS_SERVERS[:3])}{'...' if len(config.DNS_SERVERS) > 3 else ''}")
    click.echo(f"⚡ DNS rate limit: {config.DNS_RATE_LIMIT}/s")
    
    if shodan_key:
        click.echo("✅ Shodan integration enabled")
        logger.info(f"🕵️  Shodan API key provided (ends with: ...{shodan_key[-4:] if len(shodan_key) > 4 else shodan_key})")
    else:
        click.echo("⚠️  Shodan integration disabled (no API key)")
        
    if virustotal_key:
        click.echo("✅ VirusTotal integration enabled")
        logger.info(f"🛡️  VirusTotal API key provided (ends with: ...{virustotal_key[-4:] if len(virustotal_key) > 4 else virustotal_key})")
    else:
        click.echo("⚠️  VirusTotal integration disabled (no API key)")
    
    click.echo("")
    
    # Run reconnaissance
    try:
        logger.info(f"🚀 Starting reconnaissance for target: {target}")
        data = engine.run_reconnaissance(target)
        
        # Display final statistics
        stats = data.get_stats()
        click.echo("")
        click.echo("=" * 60)
        click.echo("📊 RECONNAISSANCE COMPLETE")
        click.echo("=" * 60)
        click.echo(f"🏠 Hostnames discovered: {stats['hostnames']}")
        click.echo(f"🌐 IP addresses found: {stats['ip_addresses']}")
        click.echo(f"📋 DNS records collected: {stats['dns_records']}")
        click.echo(f"📜 Certificates found: {stats['certificates']}")
        click.echo(f"🕵️  Shodan results: {stats['shodan_results']}")
        click.echo(f"🛡️  VirusTotal results: {stats['virustotal_results']}")
        
        # Calculate and display timing
        if data.end_time and data.start_time:
            duration = data.end_time - data.start_time
            click.echo(f"⏱️  Total time: {duration}")
        
        click.echo("")
        
        # Generate reports
        logger.info("📄 Generating reports...")
        report_gen = ReportGenerator(data)
        
        if output:
            # Save to files
            saved_files = []
            
            if not text_only:
                json_file = f"{output}.json"
                try:
                    json_content = data.to_json()
                    with open(json_file, 'w', encoding='utf-8') as f:
                        f.write(json_content)
                    saved_files.append(json_file)
                    logger.info(f"💾 JSON report saved: {json_file}")
                except Exception as e:
                    logger.error(f"❌ Failed to save JSON report: {e}")
            
            if not json_only:
                text_file = f"{output}.txt"
                try:
                    with open(text_file, 'w', encoding='utf-8') as f:
                        f.write(report_gen.generate_text_report())
                    saved_files.append(text_file)
                    logger.info(f"💾 Text report saved: {text_file}")
                except Exception as e:
                    logger.error(f"❌ Failed to save text report: {e}")
            
            if saved_files:
                click.echo(f"💾 Reports saved:")
                for file in saved_files:
                    click.echo(f"   📄 {file}")
        
        else:
            # Output to stdout
            if json_only:
                try:
                    click.echo(data.to_json())
                except Exception as e:
                    logger.error(f"❌ Failed to generate JSON output: {e}")
                    click.echo(f"Error generating JSON: {e}")
            elif text_only:
                try:
                    click.echo(report_gen.generate_text_report())
                except Exception as e:
                    logger.error(f"❌ Failed to generate text report: {e}")
                    click.echo(f"Error generating text report: {e}")
            else:
                # Default: show text report
                try:
                    click.echo(report_gen.generate_text_report())
                    click.echo(f"\n💡 To get JSON output, use: --json-only")
                    click.echo(f"💡 To save reports, use: --output filename")
                except Exception as e:
                    logger.error(f"❌ Failed to generate report: {e}")
                    click.echo(f"Error generating report: {e}")
    
    except KeyboardInterrupt:
        logger.warning("⚠️  Reconnaissance interrupted by user")
        click.echo("\n⚠️  Reconnaissance interrupted by user.")
        sys.exit(1)
    except Exception as e:
        logger.error(f"❌ Error during reconnaissance: {e}", exc_info=True)
        click.echo(f"❌ Error during reconnaissance: {e}")
        if verbose:
            raise  # Re-raise in verbose mode to show full traceback
        sys.exit(1)

if __name__ == '__main__':
    main()