#include "Packet.h"

#include <string>
#include <vector>
#include <locale>
#include <codecvt>
#include <thread>
#include "utils/LocalAddr.h"
#include <process.h>
#include <Tlhelp32.h>
#include <tchar.h>
#include <psapi.h>
#include <stdio.h>
#include <STDLIB.H>
#include <WtsApi32.h>
#include <UserEnv.h>
#include <iostream>

#pragma comment(lib, "iphlpapi.lib")
#pragma comment(lib, "WtsApi32.lib")
#pragma  comment (lib,"psapi")
#pragma comment(lib,"Userenv.lib")


const static int DETERMINE = 1;   //确定
const static int CANCEL = 2;   //取消

Packet::Packet() 
    : m_hbPort(0)
    , m_startTime(0)
    , m_intervalTime(0)
    , m_adhandle(nullptr)
    , m_quitFlag(false)   //退出标记
    , m_isShowToast(false)
    , m_specialPackCount(0){
}


Packet::~Packet() {
    pcap_close(m_adhandle);
    m_adhandle = nullptr;


    if (m_checkThread.joinable()) {
        m_quitFlag = true;
        m_checkThread.join();
    }

}


void Packet::Init() {
    m_hbPort = 1433;   //心跳端口
    m_startTime = time(nullptr);  //当前时间
    m_serverIp = "139.159.230.183";  //服务器ip
    m_intervalTime = 180;     //时间间隔5分钟
}


//ftp包处理  无用
void Packet::FtpPacketHandler(const struct pcap_pkthdr* header, const u_char* pkt_data) {
    bool isgetaftp = 0;
    bool isNeedOutInfo = 1;
    int head = 54;

    static std::string user;    //用户名
    static std::string password;   //密码
    static int no = 1;    //包号

    mac_header* mh = (mac_header*)pkt_data;
    ip_header* ih = (ip_header*)(pkt_data + 14);
    tcp_header* th = (tcp_header*)(pkt_data + 34);

    std::string com;
    for (int i = 0; i < 4; i++)
        com += (char)pkt_data[54 + i];

    std::string info;
    u_short sport = ntohs(th->sport);   //源端口
    u_short dport = ntohs(th->dport);   //目标端口

    if (sport == 20 || sport == 21 || dport == 20 || dport == 21)
    {
        if (header->len > 66)
        {
            if (com == "USER")
            {
                if (!user.length()) user.clear();
                isgetaftp = 1;
                isNeedOutInfo = 0;
                for (int i = head + 5; pkt_data[i] != 13; i++)
                {
                    user += (char)pkt_data[i];
                }
                info = "Input USER!";
            }
            else if (com == "PASS")
            {
                if (!password.length()) password.clear();
                isgetaftp = 1;
                isNeedOutInfo = 0;
                for (int i = head + 5; pkt_data[i] != 13; i++)
                {
                    password += (char)pkt_data[i];
                }
                info = "Input Password!";
            }
            else //if (com == "230 "|| com == "530 "||com == "220 "||com=="331 "||com=="221 ")
            {
                isgetaftp = 1;
                for (int i = head; pkt_data[i] != 13; i++)
                {
                    info += (char)pkt_data[i];
                }
            }
        }
    }

    if (isgetaftp)
    {
       std::cout << "user:" << user << "  password:" << password << std::endl;
        if (isNeedOutInfo) std::cout << "info:" << info << std::endl;
    }

    u_char flags = th->th_flags;
    if (sport == 20 || sport == 21) {
        std::cout << "SERVER to CLIENT" << std::endl;
    }
    else if (dport == 20 || dport == 21)
    {
        std::cout << "CLIENT to SERVER" << std::endl;
    }
}


bool Packet::CheckKeepAlive(u_char* option) {
    if (!option) return false;
    size_t len = strlen((const char*)option);
    std::cout << len << std::endl;
    for (int i = 0; i < len; ++i) {
        printf("%x ", *(option + i));
    }
    printf("\n");

    return len >= 4 && *option == 0x1 && *(option + 1) == 0x1 && *(option + 2) == 0x5 && *(option + 3) == 0xa;
}


void Packet::PacketHandler(u_char* param, const struct pcap_pkthdr* header, const u_char* pkt_data)
{
    Packet* pthis = (Packet*)param;
    
    //转换时间
    time_t local_tv_sec = header->ts.tv_sec;
    struct tm* ltime = localtime(&local_tv_sec);
    char timestr[16] = { 0 };
    strftime(timestr, sizeof timestr, "%H:%M:%S", ltime);

    mac_header* mh = (mac_header*)pkt_data;
    ip_header* ih = (ip_header*)(pkt_data + 14);
    tcp_header* th = (tcp_header*)(pkt_data + 34);

    std::string info;
    u_short sport = ntohs(th->sport);   //源端口
    u_short dport = ntohs(th->dport);   //目标端口

    static int no = 0;    //包号
    ++no;
    char result[1024] = { 0 };
    sprintf(result,"No.%d %s.%.6d len:%d ", no, timestr, header->ts.tv_usec, header->len);

    if (mh->type[1] == 0 && mh->type[0] == 8){//08 00    0001
        //printf("get an ip packet\n");
        if (ih->proto == 0x0006) {
            strcat_s(result, " TCP ");
        }
        else if (ih->proto == 17) {
            strcat_s(result, " UDP ");
        }
        else if (ih->proto == 1) {
            strcat_s(result, " ICMP ");
        }
        else {
            printf(" %d ", ih->proto);
            strcat_s(result, " UNKNOW ");
        }
    }

    u_char flags = th->th_flags;
    if (flags & 0x20)  strcat_s(result, "(URG)");
    if (flags & 0x10)  strcat_s(result, "(ACK)");
    if (flags & 0x08)  strcat_s(result, "(PSH)");
    if (flags & 0x04)  strcat_s(result, "(RST)");
    if (flags & 0x02)  strcat_s(result, "(SYN)");
    if (flags & 0x01)  strcat_s(result, "(FIN)");
    char srcIp[16] = { 0 }, desIp[16] = { 0 };
    sprintf_s(srcIp, "%d.%d.%d.%d", ih->saddr.byte1, ih->saddr.byte2, ih->saddr.byte3, ih->saddr.byte4);
    sprintf_s(desIp, "%d.%d.%d.%d", ih->daddr.byte1, ih->daddr.byte2, ih->daddr.byte3, ih->daddr.byte4);

    char ipInfo[64] = { 0 };
    sprintf(ipInfo," %s(%hu) -> %s(%hu)", srcIp, sport, desIp, dport); //IP和Port
    strcat_s(result, ipInfo);
    
    //char macStr[64] = { 0 };
    //strcat_s(result,"MAC ADDR:");
    //for (int i = 0; i < 6; i++) {
    //    if (i != 5){
    //        sprintf(macStr, "%02x:", mh->dest_addr[i]);
    //        strcat_s(result, macStr);
    //    }else {
    //        sprintf(macStr, "%02x", mh->dest_addr[i]);
    //        strcat_s(result, macStr);
    //    }
    //}
    //strcat_s(result," -> ");
    //for (int i = 0; i < 6; i++) {
    //    if (i != 5) {
    //        sprintf(macStr, "%02x:", mh->src_addr[i]);
    //        strcat_s(result, macStr);
    //    }else {
    //        sprintf(macStr, "%02x\n", mh->src_addr[i]);
    //        strcat_s(result, macStr);
    //    }
    //}
    //printf(result);
    

    pthis->CheckTimeOut(sport,dport,result, header->len);
}


bool Packet::IsProcessRunning(const std::wstring& processName) {
    HANDLE hProcessSnap;
    PROCESSENTRY32 pe32;

    hProcessSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (hProcessSnap == INVALID_HANDLE_VALUE) {
        return false;
    }

    pe32.dwSize = sizeof(PROCESSENTRY32);
    if (!Process32First(hProcessSnap, &pe32)) {
        CloseHandle(hProcessSnap);
        return false;
    }

    do {
        if (_wcsicmp(pe32.szExeFile, processName.c_str()) == 0) {
            CloseHandle(hProcessSnap);
            return true;
        }
    } while (Process32Next(hProcessSnap, &pe32));

    CloseHandle(hProcessSnap);
    return false;
}


void Packet::CheckTimeOut(u_short sport, u_short dport,const std::string& tcpInfoStr,int len) {
    bool isRun = IsProcessRunning(L"EdmServer.exe");

    if (!isRun)  return;    //plm进程不存在，就不打印日志

    if (dport == m_hbPort || sport == m_hbPort) {  //心跳数据
        int timeDiff = static_cast<int>(time(nullptr) - GetStartTime());
        if (timeDiff > m_intervalTime) {    //超时，调用注销脚本
            if (m_specialPackCount > 0) {
                LogoutPLM();
                LOG(DEBUG) << "Logout PLM end...";
            }
            UpdateStartTime();
        }
        else {
            static int hbCount = 0;
            ++hbCount;
            if (hbCount % 10 == 0) {  //每隔10次打印一次
                LOG(DEBUG) << tcpInfoStr << " 心跳包，" << m_intervalTime - timeDiff << "s后注销客户端...";
                hbCount = 0;
            }
        }

        //1.时间在弹窗前的最后两分钟内  m_intervalTime - timeDiff<120s
        //2.服务端 接收端口1433 长度54
        if (m_intervalTime - timeDiff <= 120 ) {
            if (dport == 1433 && 54 == len) {
                ++m_specialPackCount;
                LOG(DEBUG) << "m_specialPackCount:" << m_specialPackCount;
            }
        }
        else {
            m_specialPackCount = 0;
        }
    }
    else {
        UpdateStartTime();
        static int resetCount = 0;
        ++resetCount;
        if (resetCount % 10 == 0) {  //每隔10次打印一次
            LOG(DEBUG) << tcpInfoStr << " 其他包，重新计时...";
            resetCount = 0;
        }
    }
}


INT64 Packet::PopToast() {
    std::wstring appUserModelID = L"警告";
    WinToastLib::WinToast::instance()->setAppName(L"警告");
    WinToastLib::WinToast::instance()->setAppUserModelId(appUserModelID);
    WinToastLib::WinToastTemplate toast(WinToastLib::WinToastTemplate::ImageAndText02);
    toast.setTextField(L"即将关闭plm客户端...", WinToastLib::WinToastTemplate::FirstLine);

    char path[MAX_PATH] = { 0 };
    GetModuleFileNameA(NULL, path, MAX_PATH);

    std::string  filePath = path;
    size_t nPos = filePath.find_last_of('\\');
    filePath = filePath.substr(0, nPos);
    std::wstring_convert<std::codecvt_utf8<wchar_t>> converter;
    std::wstring wstr = converter.from_bytes(filePath + "\\notify.webp");
    toast.setImagePath(wstr);

    toast.addAction(L"确定");
    toast.addAction(L"取消");

    toast.setScenario(WinToastLib::WinToastTemplate::Scenario::Alarm);  //报警
    toast.setExpiration(10000);

    INT64 id = -1;
    if (WinToastLib::WinToast::instance()->initialize()) {
        id = WinToastLib::WinToast::instance()->showToast(toast, new MyToastHandler(this));
    }
    return id;
}

bool Packet::LogoutPLM() {
    char path[MAX_PATH] = { 0 };
    GetModuleFileNameA(NULL, path, MAX_PATH);
    std::string pathStr = path;
    size_t pos = pathStr.find_last_of('\\');
    path[pos] = '\0';
    strcat(path, "\\PDM客户端注销.BAT");
    int ret = system(path);
    LOG(DEBUG) << "Calling the logout script, ErrorCode(128ok):" << ret << ",path:" << path;
    return true;
}


void Packet::UserHandleProc(){
    char path[MAX_PATH] = { 0 };
    GetModuleFileNameA(NULL, path, MAX_PATH);

    std::string  filePath = path;
    size_t nPos = filePath.find_last_of('\\');
    filePath = filePath.substr(0, nPos);
    std::wstring_convert<std::codecvt_utf8<wchar_t>> converter;
    std::wstring uiPath = converter.from_bytes(filePath + "\\UI.exe");

    //检测线程
    m_checkThread = std::thread([this, uiPath]() {
        while (!m_quitFlag) {
            int timeDiff = static_cast<int>(time(nullptr) - GetStartTime());
            //m_specialPackCount>0表示登录状态，超时，调用注销脚本
            if (m_specialPackCount>0 && IsProcessRunning(L"EdmServer.exe")
                && timeDiff >= m_intervalTime - 10 && timeDiff <= m_intervalTime - 9) {    
                LOG(DEBUG) << "pop ui begin...";
                DWORD exit_code = -1;
                if (WaitUserHandle((LPWSTR)uiPath.data(), &exit_code)) {
                    if (exit_code == DETERMINE) {  //确定
                        LOG(DEBUG) << "点击确定...";
                        LogoutPLM();
                        UpdateStartTime();
                        Sleep(1000);
                    }
                    else if (exit_code == CANCEL) {  //取消
                        UpdateStartTime();
                        LOG(DEBUG) << "点击取消...";
                    }
                    else {
                        LOG(DEBUG) << "超时结束...";
                        LogoutPLM();
                        UpdateStartTime();
                        Sleep(1000);
                    }
                }
                else {
                    LOG(ERROR) << "UI process create failed...";
                }
            }
            Sleep(200);
        }
    });
}


int Packet::Run() {
    LOG(DEBUG) << "RUN START...";

    //char path[MAX_PATH] = { 0 };
    //GetModuleFileNameA(NULL, path, MAX_PATH);
    //std::string  filePath = path;
    //size_t nPos = filePath.find_last_of('\\');
    //filePath = filePath.substr(0, nPos);
    //std::wstring_convert<std::codecvt_utf8<wchar_t>> converter;
    //std::wstring uiPath = converter.from_bytes(filePath + "\\UI0.exe");
    //DWORD exit_code = -1;
    //WaitUserHandle((LPWSTR)uiPath.data(), &exit_code);

    UserHandleProc();
    
    char errbuf[PCAP_ERRBUF_SIZE] = { 0 };
    pcap_if_t* alldevs;
    if (pcap_findalldevs_ex((char*)PCAP_SRC_IF_STRING, NULL /* auth is not needed */, &alldevs, errbuf) == -1){
        LOG(ERROR)<< "Error in pcap_findalldevs_ex:"<< errbuf;
        return -1;
    }

    LocalAddr la;
    std::string currAdptInfo = la.GetSystemIpAddress();   //获取当前机器的IP地址
    if (currAdptInfo.empty()) {
        LOG(ERROR) << "do not get local ip...";
        return -1;
    }


    char text[2048] = {0};
    int i = 0;
    pcap_if_t* curDev = nullptr;
    for (pcap_if_t* dev = alldevs; dev != NULL; dev = dev->next) {
        char name[128] = { 0 };
        sprintf(name,"\n%d. %s ", ++i, dev->name);
        strcat_s(text, name);
        
        if (dev->description)  strcat_s(text, dev->description);
        else  strcat_s(text,"(No description available)");

        strcat_s(text, "\nIP Address: ");
        //LOG(DEBUG) << "text:"<< text;

        pcap_addr_t* addr = nullptr;
        for (addr = dev->addresses; addr != nullptr; addr = addr->next){
            if (addr->addr && addr->addr->sa_family == AF_INET){
                struct sockaddr_in* sockaddr = reinterpret_cast<struct sockaddr_in*>(addr->addr);
                char ip[24] = { 0 };
                inet_ntop(AF_INET, &(sockaddr->sin_addr), ip, sizeof(ip));
                strcat_s(text, ip);
                
                if (currAdptInfo.find(ip) != std::string::npos) {
                    curDev = dev;   //选择当前的网络适配器
                    break;
                }
            }
        }
    }
    LOG(DEBUG) << text << ",size:" << strlen(text);

    if (!curDev) {
        LOG(ERROR) << "netcord is null...";
        return -1;
    }

    int ret = 0;
    u_int netmask =0;
    do {
        if ((m_adhandle = pcap_open(curDev->name, 65536, 0 /*PCAP_OPENFLAG_PROMISCUOUS*/, 400, NULL, errbuf)) == NULL) {
            LOG(ERROR) << "Unable to open the adapter. " << curDev->name << " is not supported by WinPcap";
            ret = -1;
            break;
        }

        if (pcap_datalink(m_adhandle) != DLT_EN10MB) {
            LOG(ERROR) << "This program works only on Ethernet networks.";
            ret = -1;
            break;
        }


        if (curDev->addresses != NULL)
            netmask = ((struct sockaddr_in*)(curDev->addresses->netmask))->sin_addr.S_un.S_addr;
        else
            netmask = 0xffffff;

        struct bpf_program fcode;
        std::string packetFilter = "tcp and ip host " + m_serverIp;  //捕获过滤
        LOG(DEBUG) << "filter:" << packetFilter;
        if (pcap_compile(m_adhandle, &fcode, packetFilter.c_str(), 1, netmask) >= 0) {
            if (pcap_setfilter(m_adhandle, &fcode) < 0) {
                LOG(ERROR) << "Error setting the filter.";
                ret = -1;
                break;
            }
        }else {
            LOG(ERROR) << "Error setting the filter.";
            ret = -1;
            break;
        }
    } while (false);
  
    LOG(DEBUG) << "listening on " << curDev->description <<"...";

    pcap_freealldevs(alldevs);
    if (ret == -1)  return ret;

    pcap_loop(m_adhandle, 0, PacketHandler, (u_char*)this);
    return 0;

}
 
unsigned long Packet::ConvertMaskToULong(const std::string& maskString) {
    std::istringstream iss(maskString);
    std::string octet;
    unsigned long maskULong = 0;

    while (std::getline(iss, octet, '.')) {
        unsigned long octetULong = std::stoul(octet);
        maskULong = (maskULong << 8) | octetULong;
    }

    return maskULong;
}




BOOL Packet::WaitUserHandle(LPWSTR command_line, DWORD* exit_code)
{
    BOOL  result = FALSE;
    PROCESSENTRY32 proc_entry;

    DWORD session_id;

    HANDLE hToken = NULL;
    HANDLE hTokenDup = NULL;
    LPVOID pEnv = NULL;
    STARTUPINFO si;
    PROCESS_INFORMATION pi;

    HANDLE snap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (snap == INVALID_HANDLE_VALUE) {
        LOG(ERROR) << "CreateToolhelp32Snapshot() failed:"<< GetLastError();
        return result;
    }
    ZeroMemory(&proc_entry, sizeof(proc_entry));
    proc_entry.dwSize = sizeof(PROCESSENTRY32);
    if (!Process32First(snap, &proc_entry)) {
        LOG(ERROR) << "Process32First failed %lu" << GetLastError();
        CloseHandle(snap);
        return result;
    }
    do {
        //找出explorer进程所属用户的sessionID
        if (_tcsicmp(proc_entry.szExeFile, _T("explorer.exe")) == 0) {

            DWORD explorer_session_id = 0;
            if (ProcessIdToSessionId(proc_entry.th32ProcessID, &explorer_session_id)){
                session_id = explorer_session_id;
                break;
            }
        }
    } while (Process32Next(snap, &proc_entry));
    CloseHandle(snap);

    LOG(DEBUG) << "ui process session_id:" << session_id;

    //    根据sessionID获取Token
    if (!WTSQueryUserToken(session_id, &hToken)){
        LOG(ERROR) << "WTSQueryUserToken error:"<< GetLastError();
        CloseHandle(hToken);
        return result;
    }

    //复制新的Token
    if (!DuplicateTokenEx(hToken, MAXIMUM_ALLOWED, NULL, SecurityIdentification, TokenPrimary, &hTokenDup)){
        LOG(ERROR) << "DuplicateTokenEx error:"<< GetLastError();
        CloseHandle(hToken);
        return result;
    }

    //创建环境信息
    if (!CreateEnvironmentBlock(&pEnv, hTokenDup, FALSE)){
        LOG(ERROR) << "CreateEnvironmentBlock error："<< GetLastError();
        CloseHandle(hTokenDup);
        CloseHandle(hToken);
        return result;
    }

    //设置启动参数
    ZeroMemory(&si, sizeof(STARTUPINFO));
    si.cb = sizeof(STARTUPINFO);
    si.lpDesktop = (LPWSTR)_TEXT("winsta0\\default");

    ZeroMemory(&pi, sizeof(PROCESS_INFORMATION));

    //开始创建进程
    DWORD dwCreateFlag = NORMAL_PRIORITY_CLASS | CREATE_NEW_CONSOLE | CREATE_UNICODE_ENVIRONMENT;
    if (!CreateProcessAsUser(hToken, NULL, command_line, NULL, NULL, FALSE, dwCreateFlag, pEnv, NULL, &si, &pi)){
        LOG(ERROR) << "CreateProcessAsUser error:" << GetLastError();
    }
    else {
        result = TRUE;
    }
    
    DestroyEnvironmentBlock(pEnv);
    CloseHandle(hTokenDup);
    CloseHandle(hToken);

    int count = 100;    //最多等待10s(100*100)
    while(count--) {
        DWORD ret = WaitForSingleObject(pi.hProcess, 100);  // 等待子进程退出
        if (ret == WAIT_TIMEOUT) {
            int timeDiff = static_cast<int>(time(nullptr) - GetStartTime());
            if (timeDiff < m_intervalTime - 10) {  //用户操作了plm系统
                *exit_code = CANCEL;
                break;
            }
        }
        else if (ret == WAIT_OBJECT_0)  break;
    }

    //获取子进程的退出码
    DWORD exitCode = 0;
    GetExitCodeProcess(pi.hProcess, &exitCode);
    *exit_code = *exit_code== CANCEL ? *exit_code:exitCode;
   
    LOG(DEBUG)<<"子进程ui退出码(1:确定，2:取消):"<< *exit_code;

    // 终止子进程
    TerminateProcess(pi.hProcess, 0);
    CloseHandle(pi.hProcess);
    CloseHandle(pi.hThread);

    return result;
}