#include <iostream>
#include <map>
#include <vector>
#include <string>
#include <limits>  // Added for std::numeric_limits

// Function prototypes
void displayMainMenu();
void viewLastLoggedOn();
void viewFailedLogin();
void viewSuccessLogin();
void viewActiveSessions();
void viewRebootHistory();
void viewIPTracking();
void viewSystem();
void viewFirewall();
void viewCustomCommand();
void executeCommand(const std::string& command);

// Function to set the color
void setColor(int colorCode) {
    std::cout << "\033[" << colorCode << "m";
}

// Function to reset the color
void resetColor() {
    std::cout << "\033[0m";
}

// Function to clear the screen
void clearScreen() {
    std::cout << "\033[2J\033[1;1H";
}

// Function to print text in rainbow colors
void printRainbowText(const std::string& text) {
    int colors[] = {31, 32, 33, 34, 35, 36};  // ANSI color codes: red, green, yellow, blue, magenta, cyan
    for (size_t i = 0; i < text.length(); i++) {
        std::cout << "\033[" << colors[i % 6] << "m" << text[i];
    }
    std::cout << "\033[0m";  // Reset color
}

int main() {
    clearScreen();
    displayMainMenu();
    return 0;
}

void displayMainMenu() {
    int choice;
    do {
        printRainbowText("\nSECURITY ADMIN\n");
        std::cout << "1. Last Logged On\n";
        std::cout << "2. Failed Logins\n";
        std::cout << "3. Sign Ons\n";
        std::cout << "4. Active Sessions\n";
        std::cout << "5. Reboot History\n";
        std::cout << "6. System Resources\n";
        std::cout << "7. IP Tracking\n";
        std::cout << "8. Firewall\n";
        std::cout << "9. Custom Command\n";
        std::cout << "0. Exit\n";
        std::cout << "Select an option: ";
        std::cin >> choice;

        // Consume any leftover newline characters to avoid issues with getline later.
        std::cin.ignore(std::numeric_limits<std::streamsize>::max(), '\n');

        switch(choice) {
            case 0: 
                std::cout << "Exiting...\n";
                break;
            case 1: 
                viewLastLoggedOn(); 
                break;
            case 2: 
                viewFailedLogin(); 
                break;
            case 3: 
                viewSuccessLogin(); 
                break;
            case 4: 
                viewActiveSessions(); 
                break;
            case 5: 
                viewRebootHistory(); 
                break;
            case 6: 
                viewSystem(); 
                break;
            case 7: 
                viewIPTracking(); 
                break;
            case 8: 
                viewFirewall(); 
                break;
            case 9:
                viewCustomCommand();
                break;
            default: 
                std::cout << "Invalid choice, please try again.\n";
        }
        resetColor();
    } while(choice != 0);
}

void viewLastLoggedOn() {
    setColor(36);
    executeCommand("lastlog | grep ':'");
}

void viewFailedLogin() {
    setColor(36);
    executeCommand("cat /var/log/auth.log | grep 'Failed password'");
}

void viewSuccessLogin() {
    setColor(36);
    executeCommand("last -n 15 -i");
}

void viewActiveSessions() {
    setColor(36);
    // "w" lists currently active sessions along with system uptime and load averages.
    executeCommand("w");
}

void viewRebootHistory() {
    setColor(36);
    // "last reboot" shows the reboot history of the system.
    executeCommand("last reboot");
}

void viewSystem() {
    setColor(36);
    // Display top processes by memory and CPU usage.
    executeCommand("ps aux --sort=-%mem | head");
    executeCommand("ps aux --sort=-%cpu | head");
}

void viewIPTracking() {
    setColor(36);
    // "ip addr show" displays detailed information about network interfaces.
    executeCommand("ip addr show");
}

void viewFirewall() {
    setColor(36);
    // Check the firewall state and configuration.
    executeCommand("firewall-cmd --state");
    executeCommand("firewall-cmd --check-config");
}

void viewCustomCommand() {
    setColor(36);
    std::string customCmd;
    std::cout << "Enter custom command: ";
    std::getline(std::cin, customCmd);
    executeCommand(customCmd);
}

void executeCommand(const std::string& command) {
    std::cout << "\nExecuting: " << command << "\n\n";
    int result = system(command.c_str());
    if(result != 0) {
        std::cout << "Command failed with status " << result << "\n";
    }
}