// NullSec CppSentry - Network Packet Sentinel
// C++ security tool demonstrating:
//   - Modern C++20 features
//   - RAII and smart pointers
//   - Templates and concepts
//   - STL algorithms
//   - Structured bindings
//   - std::optional and std::variant
//
// Author: bad-antics
// License: MIT

#include <iostream>
#include <vector>
#include <map>
#include <optional>
#include <variant>
#include <algorithm>
#include <numeric>
#include <chrono>
#include <iomanip>
#include <sstream>
#include <memory>
#include <functional>

namespace nullsec {

constexpr auto VERSION = "1.0.0";

// ANSI Colors
namespace colors {
    constexpr auto RED = "\033[31m";
    constexpr auto GREEN = "\033[32m";
    constexpr auto YELLOW = "\033[33m";
    constexpr auto CYAN = "\033[36m";
    constexpr auto GRAY = "\033[90m";
    constexpr auto RESET = "\033[0m";
}

// Risk levels
enum class RiskLevel {
    Critical,
    High,
    Medium,
    Low,
    Info
};

std::string_view riskToString(RiskLevel risk) {
    switch (risk) {
        case RiskLevel::Critical: return "CRITICAL";
        case RiskLevel::High: return "HIGH";
        case RiskLevel::Medium: return "MEDIUM";
        case RiskLevel::Low: return "LOW";
        case RiskLevel::Info: return "INFO";
        default: return "UNKNOWN";
    }
}

std::string_view riskColor(RiskLevel risk) {
    switch (risk) {
        case RiskLevel::Critical:
        case RiskLevel::High:
            return colors::RED;
        case RiskLevel::Medium:
            return colors::YELLOW;
        case RiskLevel::Low:
            return colors::CYAN;
        default:
            return colors::GRAY;
    }
}

// Protocol types
enum class Protocol {
    TCP,
    UDP,
    ICMP,
    Unknown
};

std::string_view protocolToString(Protocol proto) {
    switch (proto) {
        case Protocol::TCP: return "TCP";
        case Protocol::UDP: return "UDP";
        case Protocol::ICMP: return "ICMP";
        default: return "UNKNOWN";
    }
}

// Packet structure
struct Packet {
    std::string source_ip;
    std::string dest_ip;
    uint16_t source_port;
    uint16_t dest_port;
    Protocol protocol;
    uint32_t size;
    uint8_t flags;
    std::chrono::system_clock::time_point timestamp;
    std::vector<uint8_t> payload;
};

// Detection rule
struct Rule {
    std::string id;
    std::string name;
    std::function<std::optional<RiskLevel>(const Packet&)> matcher;
    std::string description;
    std::string mitre;
};

// Alert structure
struct Alert {
    std::string rule_id;
    std::string rule_name;
    RiskLevel risk;
    Packet packet;
    std::string description;
    std::string mitre;
};

// Packet analyzer using modern C++ features
class PacketAnalyzer {
private:
    std::vector<Rule> rules_;
    std::vector<Alert> alerts_;
    std::map<std::string, int> stats_;

public:
    PacketAnalyzer() {
        initializeRules();
    }

    void initializeRules() {
        // Port scan detection
        rules_.push_back({
            "NET-001",
            "Port Scan Detected",
            [](const Packet& pkt) -> std::optional<RiskLevel> {
                // SYN flag only (scan)
                if (pkt.protocol == Protocol::TCP && (pkt.flags & 0x02) && !(pkt.flags & 0x10)) {
                    return RiskLevel::Medium;
                }
                return std::nullopt;
            },
            "TCP SYN scan activity detected",
            "T1046"
        });

        // C2 beacon detection (suspicious ports)
        rules_.push_back({
            "C2-001",
            "Suspicious C2 Port",
            [](const Packet& pkt) -> std::optional<RiskLevel> {
                static const std::vector<uint16_t> c2_ports = {4444, 5555, 6666, 31337, 8888};
                if (std::find(c2_ports.begin(), c2_ports.end(), pkt.dest_port) != c2_ports.end()) {
                    return RiskLevel::High;
                }
                return std::nullopt;
            },
            "Connection to known C2 port",
            "T1071"
        });

        // Large data transfer
        rules_.push_back({
            "EXFIL-001",
            "Large Data Transfer",
            [](const Packet& pkt) -> std::optional<RiskLevel> {
                if (pkt.size > 100000) {
                    return RiskLevel::High;
                }
                return std::nullopt;
            },
            "Unusually large data transfer detected",
            "T1048"
        });

        // ICMP tunnel detection
        rules_.push_back({
            "TUN-001",
            "ICMP Tunnel Suspected",
            [](const Packet& pkt) -> std::optional<RiskLevel> {
                if (pkt.protocol == Protocol::ICMP && pkt.size > 100) {
                    return RiskLevel::High;
                }
                return std::nullopt;
            },
            "Large ICMP packet may indicate tunneling",
            "T1095"
        });

        // DNS exfiltration
        rules_.push_back({
            "DNS-001",
            "DNS Exfiltration",
            [](const Packet& pkt) -> std::optional<RiskLevel> {
                if (pkt.protocol == Protocol::UDP && pkt.dest_port == 53 && pkt.size > 200) {
                    return RiskLevel::Critical;
                }
                return std::nullopt;
            },
            "Large DNS query may indicate data exfiltration",
            "T1048.003"
        });

        // SMB exploitation
        rules_.push_back({
            "SMB-001",
            "SMB Suspicious Activity",
            [](const Packet& pkt) -> std::optional<RiskLevel> {
                if (pkt.dest_port == 445 || pkt.dest_port == 139) {
                    return RiskLevel::Medium;
                }
                return std::nullopt;
            },
            "SMB traffic detected - potential lateral movement",
            "T1021.002"
        });

        // RDP brute force
        rules_.push_back({
            "RDP-001",
            "RDP Connection",
            [](const Packet& pkt) -> std::optional<RiskLevel> {
                if (pkt.dest_port == 3389) {
                    return RiskLevel::Low;
                }
                return std::nullopt;
            },
            "RDP connection detected",
            "T1021.001"
        });
    }

    std::vector<Alert> analyze(const Packet& packet) {
        std::vector<Alert> packet_alerts;
        
        for (const auto& rule : rules_) {
            if (auto risk = rule.matcher(packet)) {
                packet_alerts.push_back({
                    rule.id,
                    rule.name,
                    *risk,
                    packet,
                    rule.description,
                    rule.mitre
                });
                stats_[rule.id]++;
            }
        }
        
        // Store all alerts
        alerts_.insert(alerts_.end(), packet_alerts.begin(), packet_alerts.end());
        return packet_alerts;
    }

    [[nodiscard]] const std::vector<Alert>& getAlerts() const {
        return alerts_;
    }

    [[nodiscard]] const std::map<std::string, int>& getStats() const {
        return stats_;
    }
};

// Demo packet generation
std::vector<Packet> generateDemoPackets() {
    auto now = std::chrono::system_clock::now();
    
    return {
        // Normal web traffic
        {"192.168.1.100", "93.184.216.34", 54321, 443, Protocol::TCP, 1500, 0x18, now, {}},
        
        // Port scan
        {"45.33.32.156", "10.0.0.5", 12345, 22, Protocol::TCP, 64, 0x02, now, {}},
        {"45.33.32.156", "10.0.0.5", 12346, 23, Protocol::TCP, 64, 0x02, now, {}},
        {"45.33.32.156", "10.0.0.5", 12347, 80, Protocol::TCP, 64, 0x02, now, {}},
        
        // C2 communication
        {"10.0.0.50", "185.220.101.1", 49152, 4444, Protocol::TCP, 256, 0x18, now, {}},
        
        // Data exfiltration
        {"192.168.1.200", "203.0.113.50", 55555, 443, Protocol::TCP, 500000, 0x18, now, {}},
        
        // DNS exfiltration
        {"10.0.0.100", "8.8.8.8", 54000, 53, Protocol::UDP, 512, 0x00, now, {}},
        
        // ICMP tunnel
        {"192.168.1.50", "10.0.0.5", 0, 0, Protocol::ICMP, 500, 0x00, now, {}},
        
        // SMB lateral movement
        {"192.168.1.100", "192.168.1.200", 49200, 445, Protocol::TCP, 2048, 0x18, now, {}},
        
        // RDP access
        {"192.168.1.50", "10.0.0.20", 49300, 3389, Protocol::TCP, 1024, 0x18, now, {}},
        
        // Normal DNS
        {"192.168.1.100", "8.8.8.8", 54001, 53, Protocol::UDP, 64, 0x00, now, {}},
    };
}

void printBanner() {
    std::cout << "\n";
    std::cout << "╔══════════════════════════════════════════════════════════════════╗\n";
    std::cout << "║           NullSec CppSentry - Network Packet Sentinel            ║\n";
    std::cout << "╚══════════════════════════════════════════════════════════════════╝\n";
    std::cout << "\n";
}

void printUsage() {
    std::cout << "USAGE:\n";
    std::cout << "    cppsentry [OPTIONS] <INTERFACE>\n";
    std::cout << "\n";
    std::cout << "OPTIONS:\n";
    std::cout << "    -h, --help       Show this help\n";
    std::cout << "    -i, --interface  Network interface to monitor\n";
    std::cout << "    -r, --rules      Custom rules file\n";
    std::cout << "    -v, --verbose    Verbose output\n";
    std::cout << "\n";
    std::cout << "FEATURES:\n";
    std::cout << "    • Real-time packet analysis\n";
    std::cout << "    • C2/exfiltration detection\n";
    std::cout << "    • Port scan identification\n";
    std::cout << "    • MITRE ATT&CK mapping\n";
}

void printAlert(const Alert& alert) {
    std::cout << "\n";
    std::cout << "  " << riskColor(alert.risk) << "[" << riskToString(alert.risk) << "]"
              << colors::RESET << " " << alert.rule_name << "\n";
    std::cout << "    Rule:        " << alert.rule_id << "\n";
    std::cout << "    Source:      " << alert.packet.source_ip << ":" << alert.packet.source_port << "\n";
    std::cout << "    Destination: " << alert.packet.dest_ip << ":" << alert.packet.dest_port << "\n";
    std::cout << "    Protocol:    " << protocolToString(alert.packet.protocol) << "\n";
    std::cout << "    Size:        " << alert.packet.size << " bytes\n";
    std::cout << "    MITRE:       " << alert.mitre << "\n";
}

void printSummary(const std::vector<Alert>& alerts, int totalPackets) {
    auto countByRisk = [&alerts](RiskLevel risk) {
        return std::count_if(alerts.begin(), alerts.end(),
            [risk](const Alert& a) { return a.risk == risk; });
    };
    
    auto critical = countByRisk(RiskLevel::Critical);
    auto high = countByRisk(RiskLevel::High);
    auto medium = countByRisk(RiskLevel::Medium);
    auto low = countByRisk(RiskLevel::Low);
    
    std::cout << "\n";
    std::cout << colors::GRAY << "═══════════════════════════════════════════" << colors::RESET << "\n";
    std::cout << "\n";
    std::cout << "  Summary:\n";
    std::cout << "    Packets Analyzed: " << totalPackets << "\n";
    std::cout << "    Alerts Generated: " << alerts.size() << "\n";
    std::cout << "    Critical:         " << colors::RED << critical << colors::RESET << "\n";
    std::cout << "    High:             " << colors::RED << high << colors::RESET << "\n";
    std::cout << "    Medium:           " << colors::YELLOW << medium << colors::RESET << "\n";
    std::cout << "    Low:              " << colors::CYAN << low << colors::RESET << "\n";
}

void demo() {
    std::cout << colors::YELLOW << "[Demo Mode]" << colors::RESET << "\n\n";
    std::cout << colors::CYAN << "Analyzing network traffic for threats..." << colors::RESET << "\n";
    
    PacketAnalyzer analyzer;
    auto packets = generateDemoPackets();
    
    for (const auto& packet : packets) {
        analyzer.analyze(packet);
    }
    
    // Get and sort alerts
    auto alerts = analyzer.getAlerts();
    std::sort(alerts.begin(), alerts.end(),
        [](const Alert& a, const Alert& b) {
            return static_cast<int>(a.risk) < static_cast<int>(b.risk);
        });
    
    for (const auto& alert : alerts) {
        printAlert(alert);
    }
    
    printSummary(alerts, static_cast<int>(packets.size()));
}

} // namespace nullsec

int main(int argc, char* argv[]) {
    nullsec::printBanner();
    nullsec::printUsage();
    std::cout << "\n";
    nullsec::demo();
    return 0;
}