#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <netinet/in.h>
#include <libiptc/libiptc.h>
#include <linux/netfilter_ipv4/ip_tables.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#define NFC_IP_SRC_PT 0x0200
#define NFC_IP_DST_PT 0x0400

#ifndef IPT_MIN_ALIGN
#define IPT_MIN_ALIGN (__alignof__(struct ipt_entry))
#endif

#define IPT_ALIGN(s) (((s) + ((IPT_MIN_ALIGN)-1)) & ~((IPT_MIN_ALIGN)-1))

#define IPTC_ENTRY_SIZE XT_ALIGN(sizeof(struct ipt_entry))
#define IPTC_MATCH_SIZE XT_ALIGN(sizeof(struct ipt_entry_match) + sizeof(struct ipt_udp))
#define IPTC_TARGET_SIZE XT_ALIGN(sizeof(struct ipt_entry_target))
#define IPTC_FULL_SIZE IPTC_ENTRY_SIZE + IPTC_MATCH_SIZE + IPTC_TARGET_SIZE

#define IPTC_LABEL_ACCEPT "ACCEPT"
#define IPTC_LABEL_DROP "DROP"
#define IPTC_LABEL_QUEUE "QUEUE"
#define IPTC_LABEL_RETURN "RETURN"

static u_int16_t parse_port(const char *port)
{
    return atoi(port);
}

static void parse_ports(const char *portstring, u_int16_t *ports)
{
    char *buffer;
    char *cp;

    buffer = strdup(portstring);
    if ((cp = strchr(buffer, ':')) == NULL)
        ports[0] = ports[1] = parse_port(buffer);
    else {
        *cp = '\0';
        cp++;

        ports[0] = buffer[0] ? parse_port(buffer) : 0;
        ports[1] = cp[0] ? parse_port(cp) : 0xFFFF;
    }
    
    free(buffer);
}

struct ipt_entry_match *get_match(const char *sports, const char *dports, unsigned int *nfcache, const char *protocol)
{
    struct ipt_entry_match *match;
    struct ipt_udp *udpinfo;
    size_t size;

    size = IPT_ALIGN(sizeof(*match)) + IPT_ALIGN(sizeof(*udpinfo));
    match = (struct ipt_entry_match *)calloc(1, size);
    if (!match) {
        perror("calloc failed for match");
        return NULL;
    }

    match->u.match_size = size;
    strncpy(match->u.user.name, protocol, IPT_FUNCTION_MAXNAMELEN - 2);
    match->u.user.name[IPT_FUNCTION_MAXNAMELEN - 2] = '\0';

    udpinfo = (struct ipt_udp *)match->data;
    udpinfo->spts[1] = udpinfo->dpts[1] = 0xFFFF;

    if (sports) {
        *nfcache |= NFC_IP_SRC_PT;
        parse_ports(sports, udpinfo->spts);
    }
    if (dports) {
        *nfcache |= NFC_IP_DST_PT;
        parse_ports(dports, udpinfo->dpts);
    }

    return match;
}

// 打印 IP 地址
void print_ip_address(const char *label, struct in_addr addr) {
    char ip_str[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &addr, ip_str, sizeof(ip_str));
    printf("%s: %s\n", label, ip_str);
}

// 打印规则详细信息
void print_rule(const char *label, const struct ipt_entry *entry) {
    printf("%s:\n", label);
    print_ip_address("  src", entry->ip.src);
    print_ip_address("  dst", entry->ip.dst);
    printf("  proto: %u\n", entry->ip.proto);
    printf("  iniface: %s\n", entry->ip.iniface);
    printf("  outiface: %s\n", entry->ip.outiface);
    printf("  target_offset: %u\n", entry->target_offset);
    printf("  next_offset: %u\n", entry->next_offset);
    printf("  counters: pkts=%llu, bytes=%llu\n", entry->counters.pcnt, entry->counters.bcnt);
}

// 打印匹配结构详细信息
void print_match(const char *label, const struct ipt_entry_match *match) {
    printf("%s:\n", label);
    printf("  match_size: %u\n", match->u.user.match_size);
    printf("  name: %s\n", match->u.user.name);
    printf("  revision: %u\n", match->u.user.revision);
    // 如果有特定协议的数据，例如 TCP 或 UDP，可以在这里打印具体内容
    // 例如：匹配 TCP 协议时，打印源端口和目标端口
    if (strcmp(match->u.user.name, "tcp") == 0) {
        struct ipt_tcp *tcpinfo = (struct ipt_tcp *)match->data;
        printf("  TCP source port: %u\n", ntohs(tcpinfo->spts[0]));
        printf("  TCP destination port: %u\n", ntohs(tcpinfo->dpts[0]));
    } else if (strcmp(match->u.user.name, "udp") == 0) {
        struct ipt_udp *udpinfo = (struct ipt_udp *)match->data;
        printf("  UDP source port: %u\n", ntohs(udpinfo->spts[0]));
        printf("  UDP destination port: %u\n", ntohs(udpinfo->dpts[0]));
    }
}

// 打印目标结构详细信息
void print_target(const char *label, const struct ipt_entry_target *target) {
    printf("%s:\n", label);
    printf("  target_size: %u\n", target->u.user.target_size);
    printf("  name: %s\n", target->u.user.name);
    printf("  revision: %u\n", target->u.user.revision);
}

int iptc_delete_rule(const char *table, const char *chain, int protocol, const char *iniface, const char *outiface, \
    unsigned int src, unsigned int dest, const char *srcports, const char *destports, const char *target)
{
    struct xtc_handle *handle = NULL;
    struct ipt_entry *chain_entry = NULL;
    struct ipt_entry_match *entry_match = NULL;
    struct ipt_entry_target *entry_target = NULL;
    ipt_chainlabel labelit;
    long match_size = 0;
    int result = 1; // Default to failure

    // Initialize rule entry
    chain_entry = (struct ipt_entry *)calloc(1, sizeof(*chain_entry));
    if (!chain_entry) {
        perror("calloc failed");
        return 1;
    }

    if (src != 0) {
        chain_entry->ip.src.s_addr = src;
        chain_entry->ip.smsk.s_addr = 0xFFFFFFFF;
    }
    if (dest != 0) {
        chain_entry->ip.dst.s_addr = dest;
        chain_entry->ip.dmsk.s_addr = 0xFFFFFFFF;
    }

    if (iniface)
        strncpy(chain_entry->ip.iniface, iniface, IFNAMSIZ);
    if (outiface)
        strncpy(chain_entry->ip.outiface, outiface, IFNAMSIZ);

    chain_entry->ip.proto = protocol;
    if (IPPROTO_UDP == protocol)
        entry_match = get_match(srcports, destports, &chain_entry->nfcache, "udp");
    if (IPPROTO_TCP == protocol)
        entry_match = get_match(srcports, destports, &chain_entry->nfcache, "tcp");
    
    print_match("entry_match", entry_match);

    if (strcmp(target, "") == 0 || strcmp(target, IPTC_LABEL_ACCEPT) == 0 || strcmp(target, IPTC_LABEL_DROP) == 0 || strcmp(target, IPTC_LABEL_QUEUE) == 0 || strcmp(target, IPTC_LABEL_RETURN) == 0) {
        size_t size;
        size = IPT_ALIGN(sizeof(struct ipt_entry_target)) + IPT_ALIGN(sizeof(int));
        entry_target = (struct ipt_entry_target *)calloc(1, size);
        if (!entry_target) {
            perror("calloc failed for target");
            goto cleanup;
        }
        entry_target->u.user.target_size = size;
        strncpy(entry_target->u.user.name, target, XT_EXTENSION_MAXNAMELEN);
        entry_target->u.user.name[XT_EXTENSION_MAXNAMELEN - 1] = '\0';
    }
    match_size = entry_match ? entry_match->u.match_size : 0;
    print_target("entry_target", entry_target);
    
    
    struct ipt_entry *tmp_ipt = chain_entry;
    chain_entry = (struct ipt_entry *)realloc(chain_entry, sizeof(*chain_entry) + match_size + entry_target->u.target_size);
    if (!chain_entry) {
        free(tmp_ipt);
        free(entry_target);
        if (entry_match) free(entry_match);
        perror("realloc failed");
        return 1;
    }

    memcpy(chain_entry->elems + match_size, entry_target, entry_target->u.target_size);
    chain_entry->target_offset = sizeof(*chain_entry) + match_size;
    chain_entry->next_offset = sizeof(*chain_entry) + match_size + entry_target->u.target_size;
    if (entry_match) {
        memcpy(chain_entry->elems, entry_match, match_size);
    }

    handle = iptc_init(table);
    if (!handle) {
        fprintf(stderr, "libiptc error: Can't initialize table %s, %s\n", table, iptc_strerror(errno));
        goto cleanup;
    }

    strncpy(labelit, chain, sizeof(ipt_chainlabel));
    labelit[sizeof(ipt_chainlabel) - 1] = '\0';
    if (0 == iptc_is_chain(labelit, handle)) {
        fprintf(stderr, "libiptc error: Chain %s does not exist!\n", labelit);
        goto cleanup;
    }
    
    // 打印 chain_entry 详细信息
    print_rule("Chain entry", chain_entry);

    // Find and delete the matching rule
    const struct ipt_entry *entry;
    unsigned int rule_num = 0;
    for (entry = iptc_first_rule(labelit, handle); entry; (entry = iptc_next_rule(entry, handle), ++rule_num))
    {
        
        if (entry->ip.proto == chain_entry->ip.proto &&
            entry->ip.src.s_addr == chain_entry->ip.src.s_addr &&
            entry->ip.dst.s_addr == chain_entry->ip.dst.s_addr &&
                strncmp(entry->ip.iniface, chain_entry->ip.iniface, IFNAMSIZ) == 0 &&
                strncmp(entry->ip.outiface, chain_entry->ip.outiface, IFNAMSIZ) == 0)
        {
            // 打印当前 entry 详细信息
            print_rule("Current entry", entry);
            
            printf("Matching rule found. Deleting rule number: %u\n", rule_num);

            // 删除规则
            if (!iptc_delete_num_entry(labelit, rule_num, handle)) {
                fprintf(stderr, "libiptc error: Can't delete entry, %s\n", iptc_strerror(errno));
                iptc_free(handle);
                return result;
            }

            // 提交更改
            if (!iptc_commit(handle)) {
                fprintf(stderr, "libiptc error: Commit error, %s\n", iptc_strerror(errno));
            } else {
                printf("Rule deleted successfully\n");
                result = 0;
                break;
            }
        }
    }

cleanup:
    if (entry_match) free(entry_match);
    if (entry_target) free(entry_target);
    if (chain_entry) free(chain_entry);
    if (handle) iptc_free(handle);

    return result;
}

int main(int argc, char *argv[]) {
    // Test deleting rule
    unsigned int srcIp;
    inet_pton(AF_INET, argv[1], &srcIp);

    int result = iptc_delete_rule("filter", "INPUT", IPPROTO_TCP, NULL, NULL, srcIp, 0, NULL, NULL, IPTC_LABEL_DROP);

    return result;
}