#include "fksendto.h"
#include "sockqueue.h"
#include "platform.h"
#include <vector>

#ifdef _CPU_X86
#define SENDTO_ENTRY_LEN 5
#endif
#ifdef _CPU_X64
#define SENDTO_ENTRY_LEN 7
#endif

typedef int WINAPI(*sendto_func) (SOCKET, const char *, int, int, const sockaddr *, int);

static InlineHook *sendto_hook = NULL;
static sendto_func _sendto = NULL;

// 枚举当前所有可用网卡的IPv4地址
static const std::vector<in_addr> &enum_addr()
{
    static std::vector<in_addr> list;
    hostent *phost = gethostbyname(""); // 获取本机网卡
    if (phost) {
        if (phost->h_length == list.size()) // 数量相同直接返回
            return list;
        char **ppc = phost->h_addr_list; // 获取地址列表
        if (ppc) {
            list.clear();
            // 遍历列表添加到容器
            while (*ppc) {
                in_addr addr;
                memcpy(&addr, *ppc, sizeof(in_addr));
                list.push_back(addr);
                ppc++;
            }
        }
    }
    return list;
}

// hook后替换的函数
static int WINAPI fake_sendto(SOCKET s, const char *buf, int len, int flags, const sockaddr *to, int tolen)
{
    static SockQueue sockqueue;
    int result = -1;
    sockaddr_in *toaddr = (sockaddr_in *)to;
    if (toaddr->sin_addr.S_un.S_addr != INADDR_BROADCAST) {
        result = _sendto(s, buf, len, flags, to, tolen); // 非广播直接原样发送
    }
    else {
        sockaddr_in addr_self;
        int namelen = sizeof(sockaddr_in);
        getsockname(s, (sockaddr *)&addr_self, &namelen); // 获取原sockaddr
        if (addr_self.sin_port == 0) {
            // 如果没有端口号，先原样发送，这样系统才会分配一个端口号
            result = _sendto(s, buf, len, flags, to, tolen);
            getsockname(s, (sockaddr *)&addr_self, &namelen); // 重新获取
        }
        const std::vector<in_addr> &list = enum_addr();
        // 向列表中的每一个地址转发广播
        for (int i = 0; i < list.size(); i++) {
            addr_self.sin_addr = list[i]; // 把新的地址换上去，然后发送
            SOCKET sock = socket(AF_INET, SOCK_DGRAM, 0);
            BOOL opt = TRUE;
            setsockopt(sock, SOL_SOCKET, SO_BROADCAST, (char *)&opt, sizeof(BOOL)); // 广播
            setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, (char *)&opt, sizeof(BOOL)); // 重用地址端口
            bind(sock, (sockaddr *)&addr_self, sizeof(sockaddr)); // 绑定到地址端口
            result = _sendto(sock, buf, len, flags, to, tolen);
            sockqueue.add(sock); // 加到socket队列里
        }
    }
    return result;
}

void hook_sendto()
{
    if (!sendto_hook) {
        sendto_hook = new InlineHook(GetModuleHandleA("ws2_32.dll"), "sendto", (void *)fake_sendto, SENDTO_ENTRY_LEN);
        _sendto = (sendto_func)sendto_hook->get_old_entry();
        sendto_hook->hook();
    }
}

void unhook_sendto()
{
    if (sendto_hook) {
        sendto_hook->unhook();
        delete sendto_hook;
        sendto_hook = NULL;
    }
}