#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <linux/if_packet.h>
#include <net/ethernet.h>
#include <net/if.h>
#include <sys/ioctl.h>
#include <arpa/inet.h>
#include <net/if_arp.h>
#include <errno.h>
#include <time.h>

struct eth_header {
    unsigned char h_dest[ETH_ALEN];
    unsigned char h_source[ETH_ALEN];
    unsigned short h_proto;
} __attribute__((packed));

struct rarp_packet {
    unsigned short htype;
    unsigned short ptype;
    unsigned char hlen;
    unsigned char plen;
    unsigned short oper;
    unsigned char sha[ETH_ALEN];
    unsigned char spa[4];
    unsigned char tha[ETH_ALEN];
    unsigned char tpa[4];
} __attribute__((packed));

int mac_equals(unsigned char *mac1, unsigned char *mac2) {
    return memcmp(mac1, mac2, ETH_ALEN) == 0;
}

int main(int argc, char *argv[]) {
    if (argc != 2) {
        fprintf(stderr, "Usage: %s <interface_name>\n", argv[0]);
        fprintf(stderr, "Example: sudo %s eth0\n", argv[0]);
        return EXIT_FAILURE;
    }

    const char *if_name = argv[1];
    int sockfd;
    struct sockaddr_ll sa;
    struct ifreq if_idx_req, if_mac_req;
    char packet_buffer[ETH_FRAME_LEN];

    // Create raw socket
    sockfd = socket(AF_PACKET, SOCK_RAW, htons(ETH_P_RARP));
    if (sockfd == -1) {
        perror("Error creating socket. Ensure you run with sudo!");
        return EXIT_FAILURE;
    }

    // Get interface index
    memset(&if_idx_req, 0, sizeof(struct ifreq));
    strncpy(if_idx_req.ifr_name, if_name, IFNAMSIZ - 1);
    if (ioctl(sockfd, SIOCGIFINDEX, &if_idx_req) == -1) {
        perror("Error getting interface index (SIOCGIFINDEX)");
        close(sockfd);
        return EXIT_FAILURE;
    }

    // Get MAC address
    memset(&if_mac_req, 0, sizeof(struct ifreq));
    strncpy(if_mac_req.ifr_name, if_name, IFNAMSIZ - 1);
    if (ioctl(sockfd, SIOCGIFHWADDR, &if_mac_req) == -1) {
        perror("Error getting MAC address (SIOCGIFHWADDR)");
        close(sockfd);
        return EXIT_FAILURE;
    }

    // Fill Ethernet header
    struct eth_header *eth_hdr = (struct eth_header *)packet_buffer;
    memset(eth_hdr->h_dest, 0xFF, ETH_ALEN); // Broadcast
    memcpy(eth_hdr->h_source, if_mac_req.ifr_hwaddr.sa_data, ETH_ALEN);
    eth_hdr->h_proto = htons(ETH_P_RARP);

    // Fill RARP header
    struct rarp_packet *rarp_hdr = (struct rarp_packet *)(packet_buffer + sizeof(struct eth_header));
    rarp_hdr->htype = htons(ARPHRD_ETHER);
    rarp_hdr->ptype = htons(ETH_P_IP);
    rarp_hdr->hlen = ETH_ALEN;
    rarp_hdr->plen = 4;
    rarp_hdr->oper = htons(3); // ARPOP_RARP_REQUEST
    memcpy(rarp_hdr->sha, eth_hdr->h_source, ETH_ALEN);
    memset(rarp_hdr->spa, 0x00, 4);
    memcpy(rarp_hdr->tha, eth_hdr->h_source, ETH_ALEN);
    memset(rarp_hdr->tpa, 0x00, 4);

    // Set up destination socket address
    memset(&sa, 0, sizeof(struct sockaddr_ll));
    sa.sll_ifindex = if_idx_req.ifr_ifindex;
    sa.sll_halen = ETH_ALEN;
    memcpy(sa.sll_addr, eth_hdr->h_dest, ETH_ALEN);

    int packet_len = sizeof(struct eth_header) + sizeof(struct rarp_packet);
    if (sendto(sockfd, packet_buffer, packet_len, 0, (struct sockaddr *)&sa, sizeof(sa)) < 0) {
        perror("Error sending packet (sendto)");
        close(sockfd);
        return EXIT_FAILURE;
    }

    printf("RARP request sent successfully from interface %s (MAC: %02x:%02x:%02x:%02x:%02x:%02x)\n",
           if_name,
           eth_hdr->h_source[0], eth_hdr->h_source[1], eth_hdr->h_source[2],
           eth_hdr->h_source[3], eth_hdr->h_source[4], eth_hdr->h_source[5]);

    printf("Waiting for RARP reply...\n");

    // Set timeout for recv
    struct timeval timeout;
    timeout.tv_sec = 5;
    timeout.tv_usec = 0;
    setsockopt(sockfd, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));

    while (1) {
        char recv_buffer[ETH_FRAME_LEN];
        ssize_t recv_len = recvfrom(sockfd, recv_buffer, sizeof(recv_buffer), 0, NULL, NULL);
        if (recv_len < 0) {
            if (errno == EWOULDBLOCK || errno == EAGAIN) {
                printf("Timeout: No RARP reply received.\n");
            } else {
                perror("Error receiving packet");
            }
            break;
        }

        struct eth_header *recv_eth = (struct eth_header *)recv_buffer;
        if (ntohs(recv_eth->h_proto) != ETH_P_RARP)
            continue;

        struct rarp_packet *recv_rarp = (struct rarp_packet *)(recv_buffer + sizeof(struct eth_header));
        if (ntohs(recv_rarp->oper) != 4) // ARPOP_RARP_REPLY
            continue;

        if (!mac_equals(recv_rarp->tha, eth_hdr->h_source))
            continue;

        char ip_str[INET_ADDRSTRLEN];
        inet_ntop(AF_INET, recv_rarp->tpa, ip_str, sizeof(ip_str));

        printf("RARP reply received!\n");
        printf("IP address assigned: %s\n", ip_str);
        break;
    }

    close(sockfd);
    return EXIT_SUCCESS;
}