#define WIN32_LEAN_AND_MEAN

#include <stdio.h>
#include <conio.h>
#include <windows.h>
#include <fltdefs.h>
#include <winsock2.h>
#include <ws2tcpip.h>

#define DIR_INPUT    1
#define DIR_OUTPUT   2

#pragma comment(lib, "iphlpapi.lib")
#pragma comment(lib, "ws2_32.lib")

int getIpAddress(char *list)
{
    WSADATA WinsockData;
    if(WSAStartup(MAKEWORD(2, 2), &WinsockData) != 0)
        return -1;

    SOCKET sock = WSASocket(AF_INET, SOCK_DGRAM, 0, 0, 0, 0);
    if(sock == SOCKET_ERROR)
        return -1;
    
    unsigned long nBytesReturned;
    INTERFACE_INFO InterfaceList[32];
    if(WSAIoctl(sock, SIO_GET_INTERFACE_LIST, 0, 0, &InterfaceList,
        sizeof(InterfaceList), &nBytesReturned, 0, 0) == SOCKET_ERROR)
        return -1;
    
    int nNumInterfaces = nBytesReturned / sizeof(INTERFACE_INFO);
    if(0 < nNumInterfaces){
        sockaddr_in *pAddress = (sockaddr_in *) & (InterfaceList[0].iiAddress);
        lstrcpy(list, inet_ntoa(pAddress->sin_addr));
    }

    WSACleanup();
    return 0;
}

int parseAddress(char *addrStr, 
                 BYTE *addr, 
                 BYTE *mask)
{
    int ip[4], mk[4];
    int num = sscanf(addrStr, "%d.%d.%d.%d/%d.%d.%d.%d", 
        &ip[0], &ip[1], &ip[2], &ip[3],
        &mk[0], &mk[1], &mk[2], &mk[3]);
    if(num != 8)
        return -1;

    for(int i=0; i < 4; i++){
        if(ip[i] < 0 || ip[i] > 255 || mk[i] < 0 || mk[i] > 255)
            return -1;
        addr[i] = ip[i], mask[i] = mk[i];
    }
    return 0;
}

int addFilter(INTERFACE_HANDLE hIf, 
              int flag, 
              DWORD dwProtocol, 
              char *localAddrStr, 
              char *remoteAddrStr,
              WORD srcPort,
              WORD dstPort,
              WORD srcHighPort,
              WORD dstHighPort)
{
    BYTE localAddrBin[4], remoteAddrBin[4];
    BYTE localMaskBin[4], remoteMaskBin[4];
    parseAddress(localAddrStr,  localAddrBin,  localMaskBin);
    parseAddress(remoteAddrStr, remoteAddrBin, remoteMaskBin);

    PF_FILTER_DESCRIPTOR filter;
    memset(&filter, 0, sizeof(filter));
    filter.dwFilterFlags = FD_FLAGS_NOSYN;
    filter.dwRule        = 0;
    filter.pfatType      = PF_IPV4;
    filter.dwProtocol    = dwProtocol;
    filter.fLateBound    = 0;

    DWORD dwRet;
    if(flag == DIR_INPUT){
        filter.SrcAddr = remoteAddrBin, filter.SrcMask = remoteMaskBin;
        filter.DstAddr = localAddrBin,  filter.DstMask = localMaskBin;
        filter.wSrcPort = dstPort, filter.wDstPort = srcPort;
        filter.wSrcPortHighRange = dstHighPort;
        filter.wDstPortHighRange = srcHighPort;
        dwRet = PfAddFiltersToInterface(hIf, 1, &filter, 0, NULL, NULL);
    }else{
        filter.SrcAddr = localAddrBin, filter.SrcMask = localMaskBin;
        filter.DstAddr = remoteAddrBin, filter.DstMask = remoteMaskBin;
        filter.wSrcPort = srcPort, filter.wDstPort = dstPort;
        filter.wSrcPortHighRange = srcHighPort;
        filter.wDstPortHighRange = dstHighPort;
        dwRet = PfAddFiltersToInterface(hIf, 0, NULL, 1, &filter, NULL);
    }

    if(dwRet != NO_ERROR){
        printf("addFilter: PfAddFiltersToInterface failed");
        return -1;
    }

    return 0;
}

int main(void)
{
    char useip[32] = "192.168.4.2";

    char iplist[32];
    if(!getIpAddress(iplist))
        lstrcpy(useip, iplist);

    INTERFACE_HANDLE hFilterIf;
    DWORD dwRet = PfCreateInterface(
        NULL, PF_ACTION_FORWARD, PF_ACTION_FORWARD, FALSE, TRUE, &hFilterIf);
    if(dwRet != NO_ERROR){
        fprintf(stderr, "PfCreateInterface failed\r\n");
        return -1;
    }

    lstrcat(useip, "/255.255.255.255");
    
    addFilter(hFilterIf, DIR_OUTPUT, FILTER_PROTO_TCP,
        useip, "0.0.0.0/0.0.0.0", FILTER_TCPUDP_PORT_ANY, 20,
        FILTER_TCPUDP_PORT_ANY, 25);
    addFilter(hFilterIf, DIR_INPUT,  FILTER_PROTO_TCP,
        useip, "0.0.0.0/0.0.0.0", FILTER_TCPUDP_PORT_ANY, 20,
        FILTER_TCPUDP_PORT_ANY, 25);

    addFilter(hFilterIf, DIR_OUTPUT, FILTER_PROTO_TCP,
        useip, "0.0.0.0/0.0.0.0", FILTER_TCPUDP_PORT_ANY, 81,
        FILTER_TCPUDP_PORT_ANY, 65535);
    addFilter(hFilterIf, DIR_INPUT,  FILTER_PROTO_TCP,
        useip, "0.0.0.0/0.0.0.0", FILTER_TCPUDP_PORT_ANY, 81,
        FILTER_TCPUDP_PORT_ANY, 65535);

    int ip[4], mk[4];
    int num = sscanf(useip, "%d.%d.%d.%d/%d.%d.%d.%d",
        &ip[0], &ip[1], &ip[2], &ip[3],
        &mk[0], &mk[1], &mk[2], &mk[3]);
    if(num != 8)
        return -1;

    BYTE localAddr[4];
    for(int i=0; i < 4; i++)
        localAddr[i] = ip[i];

    printf("Start Filtering");

    dwRet = PfBindInterfaceToIPAddress(hFilterIf, PF_IPV4, (PBYTE)localAddr);
    if(dwRet != NO_ERROR){
        fprintf(stderr, "PfBindInterfaceToIPAddress failed");
        _getch();
        return -1;
    }

    _getch();

    PfUnBindInterface(hFilterIf);
    PfDeleteInterface(hFilterIf);
    return 0;
}