#include <dlfcn.h>
#include <fcntl.h>
#include <cstring>
#include <string>
#include <set>
#include <unistd.h>
#include <algorithm>
#include <arpa/inet.h>

#define ACCEPT 9999
#define CONNECT 19999


// 定义指向原始文件访问函数的函数指针
typedef int (*file_access_func_t)(const char *pathname, int mode, ...);
static file_access_func_t original_file_access = nullptr;
// 定义指向原始open函数的函数指针
typedef int (*file_open_func_t)(const char *pathname, int flags, ...);
static file_open_func_t original_file_open = nullptr;
// 定义指向原始creat函数的函数指针
typedef int (*file_creat_func_t)(const char *pathname, mode_t mode);
static file_creat_func_t original_file_creat = nullptr;
//read
typedef ssize_t (*file_read_func_t)(int fd, void *buf, size_t nbytes);
static file_read_func_t original_file_read = nullptr;
// connect
typedef int (*ip_connect_fun_t)(int, const struct sockaddr *, socklen_t);
static ip_connect_fun_t original_file_connect = nullptr;
// accept
typedef int (*ip_accept_fun_t)(int, struct sockaddr *, socklen_t *);
static ip_accept_fun_t original_file_accept = nullptr;


// 定义原始的do_sys_open()函数指针
//static long (*real_do_sys_open)(int dfd, const char __user *filename, int flags, umode_t mode) = nullptr;

// 定义日志文件路径和配置文件路径的环境变量
const char *LOG_FILE_PATH = "/home/dongl/code/safe_test/log";
const char *CONFIG_FILE_PATH = "/home/dongl/code/safe_test/config.txt";
const char *LOG_NET_IP_PATH = "/home/dongl/code/safe_test/net-log";
const char *CONFIG_NET_IP_PATH = "/home/dongl/code/safe_test/net-config.txt";

// 将时间戳转换为字符串
std::string Time_t2String(time_t t) {
    char tmp[64] = {0};
    strftime(tmp, sizeof(tmp), "%Y-%m-%d %H:%M:%S", localtime(&t));
    return tmp;
}

// 获取当前进程的可执行文件路径和名称
size_t get_executable_path(char *processdir, char *processname, size_t len) {
    char *path_end;
    if (readlink("/proc/self/exe", processdir, len) <= 0)
        return -1;
    path_end = strrchr(processdir, '/');
    if (path_end == nullptr)
        return -1;
    ++path_end;
    strcpy(processname, path_end);
    *path_end = '\0';
    return (size_t) (path_end - processdir);
}

// 读取配置文件并填充文件和进程集合
void read_config_file(std::set<std::string>& file_set, std::set<std::string>& process_set, const char* config_name) {
    // 打开配置文件
    int config_fd = original_file_open(config_name, O_RDONLY);
    if (config_fd < 0)
        return;

    // 读取配置文件
    char buff[4096];
    std::string data;
    bool state = true;
    while (original_file_read(config_fd, buff, sizeof(buff)) > 0) {
        for (const auto &item: buff) {
            if (item != '\n') {
                data.push_back(item);
            } else if (item == '\n') {
                // 检查行是文件部分还是进程部分
                std::string temp = config_name == CONFIG_FILE_PATH ? "[safe_file]" : "[black_ip]";
                if (data == temp) {
                    state = true;
                    data.clear();
                    continue;
                } else if (data == "[safe_process]") {
                    state = false;
                    data.clear();
                    continue;
                }

                // 将文件或进程添加到相应的集合中
                if (state) {
                    file_set.insert(data);
                    data.clear();
                } else {
                    process_set.insert(data);
                    data.clear();
                }
            }
        }
    }

    // 关闭配置文件
    close(config_fd);
}

// 将文件访问日志写入日志文件
void write_log(const std::string &pathname, int mode, const std::string &result, const char* log_name) {
    char run_path[40960] = {0};
    char process_name[10240] = {0};
    get_executable_path(run_path, process_name, sizeof(run_path));

    std::string access_mode;
    if (mode == O_RDONLY)
        access_mode = "读取";
    else if (mode == ACCEPT)
        access_mode = "ACCEPT";
    else if (mode == CONNECT)
        access_mode = "CONNECT";
    else if (mode == O_WRONLY)
        access_mode = "写入";
    else if (mode == O_RDWR)
        access_mode = "读写";
    else if (mode == (O_WRONLY | O_APPEND))
        access_mode = "追加";

    std::string log = "[" + Time_t2String(time(nullptr)) + "] " +
                      "<" + pathname + "> " + access_mode + " {" + process_name + ":" + run_path + "} " + result + "\n";
    int log_fd = original_file_open(log_name, O_WRONLY | O_CREAT | O_APPEND, 0644);
    if (log_fd > 0) {
        write(log_fd, log.c_str(), log.size());
        close(log_fd);
    }
}

//
ssize_t read(int fd, void *buf, size_t nbytes) {
    std::set<std::string> file_set;
    std::set<std::string> process_set;
    read_config_file(file_set, process_set, CONFIG_FILE_PATH);
    // 查进程名 及 运行路径
    char run_path[4096] = {0};
    char process_name[1024] = {0};
    get_executable_path(run_path, process_name, sizeof(run_path));
    // 根据fd取文件名
    char file_name[4096] = {'\0'};
    char file_path[4096] = {'\0'};
    snprintf(file_name, sizeof(file_name), "/proc/%d/fd/%d", getpid(), fd);
    ssize_t ret = readlink(file_name, file_path, sizeof(file_path) - 1);
    if (ret == -1) {
        return original_file_read(fd, buf, nbytes);
    }

    // 检查
    auto safe_file = file_set.find(file_path);
    auto safe_process = process_set.find(process_name);

    // 校验
    if (safe_file != file_set.end()) {
        if (safe_process != process_set.end()) {
            write_log(file_path, O_RDONLY, "正常访问", LOG_FILE_PATH);
        } else {
            write_log(file_path, O_RDONLY, "拒绝访问", LOG_FILE_PATH);
            errno = EACCES;
            return original_file_read(fd, buf, 0);
        }
    }
    return original_file_read(fd, buf, nbytes);
}

// 判断接入的网络是否为过滤IP，如果是则提前返回
int connect(int sockfd, const struct sockaddr *addr, socklen_t addrlen) {
    // 获取链接ip
    char ip[128];
    memset(ip, 0, sizeof(ip));  // ip地址
    int port = -1;              // 端口号
    if (AF_INET == addr->sa_family) {
        auto *sa4 = (struct sockaddr_in *) addr;
        inet_ntop(AF_INET, (void *) (struct sockaddr *) &sa4->sin_addr, ip, sizeof(ip));
        port = ntohs(sa4->sin_port);
        printf("\nAF_INET IP===%s:%d\n", ip, port);
    }

    // 查进程名 及 运行路径
    char run_path[4096] = {0};
    char process_name[1024] = {0};
    get_executable_path(run_path, process_name, sizeof(run_path));

    // 获取黑名单ip
    std::set<std::string> ip_set;
    std::set<std::string> process_set;
    read_config_file(ip_set, process_set, CONFIG_NET_IP_PATH);
    auto black_ip = ip_set.find(ip);
    auto safe_process = process_set.find(process_name);

    // 比较IP  是黑名单ip
    if (black_ip != ip_set.end()) {
        // log ip
        std::string temp;
        temp.append(ip);
        temp.append(":"+std::to_string(port)+"<"+std::to_string(sockfd)+">");

        if (safe_process != process_set.end()) {
            write_log(temp,  CONNECT, "允许访问", LOG_NET_IP_PATH);
            return original_file_connect(sockfd, addr, addrlen);
        } else {
            write_log(temp,  CONNECT, "拒绝访问", LOG_NET_IP_PATH);
            return -1;
        }
    }

    return original_file_connect(sockfd, addr, addrlen);
}

int accept(int sockfd, struct sockaddr *addr, socklen_t *addrlen) {
    // 获取链接ip
    char ip[128];
    memset(ip, 0, sizeof(ip));  // ip地址
    int port = -1;              // 端口号
    if (AF_INET == addr->sa_family) {
        auto *sa4 = (struct sockaddr_in *) addr;
        inet_ntop(AF_INET, (void *) (struct sockaddr *) &sa4->sin_addr, ip, sizeof(ip));
        port = ntohs(sa4->sin_port);
        printf("\nAF_INET IP===%s:%d\n", ip, port);
    }

    // 查进程名 及 运行路径
    char run_path[4096] = {0};
    char process_name[1024] = {0};
    get_executable_path(run_path, process_name, sizeof(run_path));

    // 获取黑名单ip
    std::set<std::string> ip_set;
    std::set<std::string> process_set;
    read_config_file(ip_set, process_set, CONFIG_NET_IP_PATH);
    auto black_ip = ip_set.find(ip);
    auto safe_process = process_set.find(process_name);

    // log ip
    std::string temp;
    temp.append(ip);
    temp.append(":"+std::to_string(port)+"<"+std::to_string(sockfd)+">");

    // 比较IP  是黑名单ip
    if (black_ip != ip_set.end()) {
        if (safe_process != process_set.end()) {
            write_log(temp,  ACCEPT, "允许访问",LOG_NET_IP_PATH);
            return original_file_accept(sockfd, addr, addrlen);
        } else {
            write_log(temp,  ACCEPT, "拒绝访问", LOG_NET_IP_PATH);
            return -1;
        }
    }

    // 放行
    return original_file_accept(sockfd, addr, addrlen);
}


// 初始化函数，获取原始文件访问函数的地址并替换为拦截函数
__attribute__((constructor)) void init() {
    original_file_access = (file_access_func_t) dlsym(RTLD_NEXT, "access");
    original_file_open = (file_open_func_t) dlsym(RTLD_NEXT, "open");
    original_file_creat = (file_creat_func_t) dlsym(RTLD_NEXT, "creat");
    original_file_read = (file_read_func_t) dlsym(RTLD_NEXT, "read");
    original_file_connect = (ip_connect_fun_t) dlsym(RTLD_NEXT, "connect");
    original_file_accept = (ip_accept_fun_t) dlsym(RTLD_NEXT, "accept");
}