/***************************************************************
 * 文件：resolver.c
 * 功能：C 語言實現(xiàn)域名解析（兩種方式）
 *   - 方法 A：使用系統(tǒng)接口 getaddrinfo/getnameinfo（推薦）
 *   - 方法 B：自實現(xiàn)簡易 DNS 客戶端（UDP 查詢 A / AAAA / CNAME）
 *
 * 編譯（Linux）:
 *   gcc resolver.c -o resolver
 *
 * 用法示例：
 *   ./resolver sys www.example.com         # 使用系統(tǒng)接口解析
 *   ./resolver dns 8.8.8.8 www.example.com # 使用自定義 DNS 客戶端，指定 DNS 服務(wù)器
 *   ./resolver reverse sys 93.184.216.34   # 反向解析（IP->域名）使用系統(tǒng)接口
 *
 ***************************************************************/

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>

#include <unistd.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <time.h>
#include <sys/time.h>

/* -----------------------------
   通用工具函數(shù)與定義
   ----------------------------- */

#define MAXBUF 512

static void pprint_addrinfo(struct addrinfo *res) {
    char host[NI_MAXHOST];
    for (struct addrinfo *p = res; p != NULL; p = p->ai_next) {
        void *addr;
        const char *ipver;
        if (p->ai_family == AF_INET) {
            struct sockaddr_in *ipv4 = (struct sockaddr_in *)p->ai_addr;
            addr = &(ipv4->sin_addr);
            ipver = "IPv4";
        } else if (p->ai_family == AF_INET6) {
            struct sockaddr_in6 *ipv6 = (struct sockaddr_in6 *)p->ai_addr;
            addr = &(ipv6->sin6_addr);
            ipver = "IPv6";
        } else {
            continue;
        }
        if (inet_ntop(p->ai_family, addr, host, sizeof(host)) == NULL) {
            strncpy(host, "unknown", sizeof(host));
            host[sizeof(host)-1] = 0;
        }
        printf("  %-4s  %s\n", ipver, host);
    }
}

/* -----------------------------
   方法 A：使用系統(tǒng)接口 getaddrinfo / getnameinfo
   ----------------------------- */

static int resolve_with_getaddrinfo(const char *name) {
    struct addrinfo hints, *res;
    int rv;

    memset(&hints, 0, sizeof(hints));
    hints.ai_family = AF_UNSPEC;    // 支持 IPv4 和 IPv6
    hints.ai_socktype = 0;          // 任意 socket 類型
    hints.ai_flags = AI_CANONNAME;  // 返回規(guī)范名（如可用）

    rv = getaddrinfo(name, NULL, &hints, &res);
    if (rv != 0) {
        fprintf(stderr, "getaddrinfo: %s\n", gai_strerror(rv));
        return 1;
    }

    printf("Canonical name: %s\n", res->ai_canonname ? res->ai_canonname : "(none)");
    printf("Addresses for %s:\n", name);
    pprint_addrinfo(res);

    freeaddrinfo(res);
    return 0;
}

/* 反向解析：IP -> 主機名 */
static int reverse_with_getnameinfo(const char *ipstr) {
    struct sockaddr_storage sa;
    socklen_t sa_len;
    char host[NI_MAXHOST];

    memset(&sa, 0, sizeof(sa));
    if (strchr(ipstr, ':')) {
        // IPv6
        struct sockaddr_in6 *sa6 = (struct sockaddr_in6 *)&sa;
        sa6->sin6_family = AF_INET6;
        if (inet_pton(AF_INET6, ipstr, &(sa6->sin6_addr)) != 1) {
            fprintf(stderr, "Invalid IPv6 address: %s\n", ipstr);
            return 1;
        }
        sa_len = sizeof(struct sockaddr_in6);
    } else {
        // IPv4
        struct sockaddr_in *sa4 = (struct sockaddr_in *)&sa;
        sa4->sin_family = AF_INET;
        if (inet_pton(AF_INET, ipstr, &(sa4->sin_addr)) != 1) {
            fprintf(stderr, "Invalid IPv4 address: %s\n", ipstr);
            return 1;
        }
        sa_len = sizeof(struct sockaddr_in);
    }

    int rv = getnameinfo((struct sockaddr *)&sa, sa_len, host, sizeof(host), NULL, 0, NI_NAMEREQD);
    if (rv != 0) {
        fprintf(stderr, "getnameinfo: %s\n", gai_strerror(rv));
        return 1;
    }
    printf("Reverse lookup: %s -> %s\n", ipstr, host);
    return 0;
}

/* -----------------------------
   方法 B：自實現(xiàn)簡易 DNS 客戶端（UDP）
   僅支持查詢 A(1) / AAAA(28) / CNAME (5)
   ----------------------------- */

/* DNS header: 12 bytes */
#pragma pack(push, 1)
struct dns_header {
    unsigned short id;
    unsigned short flags;
    unsigned short qdcount;
    unsigned short ancount;
    unsigned short nscount;
    unsigned short arcount;
};
#pragma pack(pop)

/* DNS question tail (type & class) */
struct dns_question_tail {
    unsigned short qtype;
    unsigned short qclass;
};

/* 資源記錄固定頭部（不含可變長度 NAME 與 RDATA） */
#pragma pack(push,1)
struct dns_rr_fixed {
    unsigned short type;
    unsigned short class;
    unsigned int ttl;
    unsigned short rdlength;
};
#pragma pack(pop)

/* 將域名從點分格式轉(zhuǎn)換為 DNS 報文的標(biāo)簽格式：
   e.g., "www.example.com" -> [3]www[7]example[3]com[0]
   返回寫入的字節(jié)數(shù)到 buf（不超過 buflen）。
*/
static int dns_name_pack(const char *name, unsigned char *buf, int buflen) {
    int nlen = strlen(name);
    if (nlen == 0) {
        if (buflen < 1) return -1;
        buf[0] = 0;
        return 1;
    }
    int pos = 0;
    const char *label = name;
    const char *p = name;
    while (1) {
        if (*p == '.' || *p == '\0') {
            int len = p - label;
            if (len > 63) return -1; // label too long
            if (pos + 1 + len >= buflen) return -1;
            buf[pos++] = (unsigned char)len;
            if (len > 0) {
                memcpy(&buf[pos], label, len);
                pos += len;
            }
            if (*p == '\0') {
                // terminate with zero
                if (pos >= buflen) return -1;
                buf[pos++] = 0;
                break;
            }
            label = p + 1;
        }
        p++;
    }
    return pos;
}

/* 解析 DNS 報文中的名稱（包含指針壓縮）
   packet：完整報文起始地址
   pktlen：報文總長度
   offset：當(dāng)前讀取偏移（輸入時指向名稱開始），函數(shù)結(jié)束時 offset 會更新到名稱后的第一個字節(jié)
   out：輸出緩沖區(qū)存放解析出的點分格式字符串
   outlen：輸出長度
   返回：0 成功，-1 失敗
*/
static int dns_name_unpack(unsigned char *packet, int pktlen, int *offset, char *out, int outlen) {
    int orig_offset = *offset;
    int pos = 0;
    int jumped = 0;
    int max_jumps = 0;
    int cur = *offset;
    while (cur < pktlen) {
        unsigned char len = packet[cur];
        if (len == 0) {
            // end of name
            if (!jumped) *offset = cur + 1;
            if (pos == 0) {
                // root
                if (pos + 1 > outlen) return -1;
                out[pos] = '\0';
            } else {
                if (pos + 1 > outlen) return -1;
                out[pos] = '\0';
            }
            return 0;
        }
        if ((len & 0xC0) == 0xC0) {
            // pointer: next byte + (len & 0x3F) << 8
            if (cur + 1 >= pktlen) return -1;
            int b2 = packet[cur + 1];
            int pointer = ((len & 0x3F) << 8) | b2;
            if (pointer >= pktlen) return -1;
            if (!jumped) *offset = cur + 2;
            cur = pointer;
            jumped = 1;
            if (++max_jumps > 10) return -1; // avoid loops
            continue;
        } else {
            // label
            cur++;
            if (cur + len > pktlen) return -1;
            if (pos + len + 1 >= outlen) return -1;
            memcpy(out + pos, packet + cur, len);
            pos += len;
            out[pos++] = '.';
            cur += len;
        }
    }
    return -1;
}

/* 構(gòu)造并發(fā)送 DNS 查詢（type: 1=A, 28=AAAA）
   dns_server: dotted IP string, e.g., "8.8.8.8"
   name: 域名
   type: qtype (1/A, 28/AAAA)
   timeout_sec: 接收超時時間（秒）
   返回：0 成功（并將打印解析結(jié)果），非0表示失敗
*/
static int dns_query_simple(const char *dns_server, const char *name, unsigned short qtype, int timeout_sec) {
    unsigned char buf[512];
    memset(buf, 0, sizeof(buf));

    struct dns_header hdr;
    memset(&hdr, 0, sizeof(hdr));
    srand((unsigned int)time(NULL));
    hdr.id = (unsigned short) (rand() & 0xFFFF);
    hdr.flags = htons(0x0100); // standard query, recursion desired
    hdr.qdcount = htons(1);

    // 寫 header
    memcpy(buf, &hdr, sizeof(hdr));
    int offset = sizeof(hdr);

    // 寫 QNAME
    int n = dns_name_pack(name, buf + offset, sizeof(buf) - offset);
    if (n < 0) {
        fprintf(stderr, "dns_name_pack failed\n");
        return 1;
    }
    offset += n;

    // 寫 QTYPE & QCLASS
    struct dns_question_tail qt;
    qt.qtype = htons(qtype);
    qt.qclass = htons(1); // IN
    memcpy(buf + offset, &qt, sizeof(qt));
    offset += sizeof(qt);

    // 發(fā)送 UDP 包
    int sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock < 0) {
        perror("socket");
        return 1;
    }

    struct sockaddr_in serv;
    memset(&serv, 0, sizeof(serv));
    serv.sin_family = AF_INET;
    serv.sin_port = htons(53);
    if (inet_pton(AF_INET, dns_server, &serv.sin_addr) != 1) {
        fprintf(stderr, "Invalid DNS server IP: %s\n", dns_server);
        close(sock);
        return 1;
    }

    // 發(fā)送并等待回復(fù)（帶重試）
    int tries = 3;
    int rv = 1;
    for (int t = 0; t < tries; t++) {
        ssize_t sent = sendto(sock, buf, offset, 0, (struct sockaddr*)&serv, sizeof(serv));
        if (sent != offset) {
            perror("sendto");
            continue;
        }

        // 設(shè)置 recv 超時
        struct timeval tv;
        tv.tv_sec = timeout_sec;
        tv.tv_usec = 0;
        setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, (const char*)&tv, sizeof(tv));

        unsigned char resp[4096];
        struct sockaddr_in from;
        socklen_t fromlen = sizeof(from);
        ssize_t rlen = recvfrom(sock, resp, sizeof(resp), 0, (struct sockaddr*)&from, &fromlen);
        if (rlen < 0) {
            if (errno == EAGAIN || errno == EWOULDBLOCK) {
                // timeout
                if (t == tries - 1) {
                    fprintf(stderr, "DNS query timeout\n");
                } else {
                    // retry
                    continue;
                }
            } else {
                perror("recvfrom");
            }
            continue;
        }

        // parse response
        if (rlen < (ssize_t)sizeof(struct dns_header)) {
            fprintf(stderr, "DNS response too short\n");
            continue;
        }
        struct dns_header rhdr;
        memcpy(&rhdr, resp, sizeof(rhdr));
        unsigned short qdcount = ntohs(rhdr.qdcount);
        unsigned short ancount = ntohs(rhdr.ancount);
        // skip questions
        int roff = sizeof(struct dns_header);
        for (int i = 0; i < qdcount; i++) {
            char qname[256];
            if (dns_name_unpack(resp, rlen, &roff, qname, sizeof(qname)) != 0) {
                fprintf(stderr, "Failed to unpack question name\n");
                break;
            }
            // skip qtype & qclass
            if (roff + sizeof(struct dns_question_tail) > rlen) break;
            roff += sizeof(struct dns_question_tail);
        }

        // parse answers
        printf("Answers (%d):\n", ancount);
        for (int i = 0; i < ancount; i++) {
            char aname[512];
            if (dns_name_unpack(resp, rlen, &roff, aname, sizeof(aname)) != 0) {
                fprintf(stderr, "Failed to unpack answer name\n");
                break;
            }
            if (roff + sizeof(struct dns_rr_fixed) > rlen) {
                fprintf(stderr, "Truncated RR header\n");
                break;
            }
            struct dns_rr_fixed rrf;
            memcpy(&rrf, resp + roff, sizeof(rrf));
            roff += sizeof(rrf);
            unsigned short type = ntohs(rrf.type);
            unsigned short rclass = ntohs(rrf.class);
            unsigned short rdlen = ntohs(rrf.rdlength);
            if (roff + rdlen > rlen) {
                fprintf(stderr, "Truncated RDATA\n");
                break;
            }
            if (type == 1 && rdlen == 4) {
                // A
                char addrbuf[INET_ADDRSTRLEN];
                inet_ntop(AF_INET, resp + roff, addrbuf, sizeof(addrbuf));
                printf("  NAME: %s  TYPE=A  ADDR=%s\n", aname, addrbuf);
            } else if (type == 28 && rdlen == 16) {
                // AAAA
                char addrbuf[INET6_ADDRSTRLEN];
                inet_ntop(AF_INET6, resp + roff, addrbuf, sizeof(addrbuf));
                printf("  NAME: %s  TYPE=AAAA ADDR=%s\n", aname, addrbuf);
            } else if (type == 5) {
                // CNAME (rdata is a name)
                int tmp_off = roff;
                char cname[512];
                if (dns_name_unpack(resp, rlen, &tmp_off, cname, sizeof(cname)) == 0) {
                    printf("  NAME: %s  TYPE=CNAME RDATA=%s\n", aname, cname);
                } else {
                    printf("  NAME: %s  TYPE=CNAME (unpack failed)\n", aname);
                }
            } else {
                printf("  NAME: %s  TYPE=%d  (rdlen=%d bytes)\n", aname, type, rdlen);
            }
            roff += rdlen;
        }

        rv = 0; // success
        break;
    }

    close(sock);
    return rv;
}

/* -----------------------------
   簡單命令行接口
   支持:
     resolver sys <domain>
     resolver reverse sys <ip>
     resolver dns <dns_server> <domain>
   ----------------------------- */

int main(int argc, char *argv[]) {
    if (argc < 3) {
        fprintf(stderr, "Usage:\n");
        fprintf(stderr, "  %s sys <domain>                  # use system resolver\n", argv[0]);
        fprintf(stderr, "  %s reverse sys <ip>              # reverse lookup via system resolver\n", argv[0]);
        fprintf(stderr, "  %s dns <dns_server> <domain>     # use custom DNS client\n", argv[0]);
        return 1;
    }

    if (strcmp(argv[1], "sys") == 0) {
        // system resolver
        return resolve_with_getaddrinfo(argv[2]);
    } else if (strcmp(argv[1], "reverse") == 0 && argc >= 4 && strcmp(argv[2], "sys") == 0) {
        return reverse_with_getnameinfo(argv[3]);
    } else if (strcmp(argv[1], "dns") == 0 && argc >= 4) {
        const char *dns_server = argv[2];
        const char *domain = argv[3];
        // query A and AAAA
        printf("Query A records via DNS server %s for %s\n", dns_server, domain);
        dns_query_simple(dns_server, domain, 1, 2);
        printf("\nQuery AAAA records via DNS server %s for %s\n", dns_server, domain);
        dns_query_simple(dns_server, domain, 28, 2);
        return 0;
    } else {
        fprintf(stderr, "Unknown command or insufficient args\n");
        return 1;
    }
}