PLMTool/main.cpp

#include "Packet.h"

#include <iostream>
#include <stdio.h>
#include <fstream>
#include <Windows.h>
#include <string>
#include <tlhelp32.h>
#include <sstream>
#include <tchar.h>
#include "easylog/easylogging++.h"

INITIALIZE_EASYLOGGINGPP


// Windows服务的服务名
TCHAR SERVICE_NAME[] = _T("MyPcap");
SERVICE_STATUS_HANDLE g_ServiceStatusHandle;
SERVICE_STATUS g_ServiceStatus;
DWORD g_ThreadID;
std::wofstream g_ofs;

// 服务主函数
VOID WINAPI ServiceMain(DWORD argc, LPWSTR* argv);

// 服务控制处理函数
VOID WINAPI ServiceCtrlHandler(DWORD CtrlCode);


void EasyLogConf()
{
    //加载默认配置  init log
    el::Configurations conf;
    //设置为默认
    conf.setToDefault();
    //设置日志输出格式
    conf.setGlobally(el::ConfigurationType::Format, "%datetime %level [%fbase|%line] %msg");
    //设置日志文件目录以及文件名
    conf.setGlobally(el::ConfigurationType::Filename, "logs\\%datetime{%Y%M%d}.log");
    //启用日志
    conf.setGlobally(el::ConfigurationType::Enabled, "true");
    //是否写入文件
    conf.setGlobally(el::ConfigurationType::ToFile, "true");
    //是否输出控制台
    conf.setGlobally(el::ConfigurationType::ToStandardOutput, "false");

    //设置配置文件
    el::Loggers::reconfigureAllLoggers(conf);    //加载配置文件
}


std::wstring GetStrNow()
{
    wchar_t buffer[128];
    SYSTEMTIME snow;
    GetLocalTime(&snow);
    swprintf_s(buffer, sizeof(buffer), L"%04d-%02d-%02d %02d:%02d:%02d", snow.wYear, snow.wMonth, snow.wDay, snow.wHour, snow.wMinute, snow.wSecond);
    return std::wstring(buffer);
}


void LogEvent(LPTSTR chmsg) {
    HANDLE hEventSource;
    LPTSTR lpszStrings[1];
    
    lpszStrings[0] = chmsg;
    
    hEventSource = RegisterEventSource(NULL, SERVICE_NAME);
    if (hEventSource != NULL)
    {
        ReportEvent(hEventSource, EVENTLOG_INFORMATION_TYPE, 0, 0, NULL, 1, 0, (LPCTSTR*)&lpszStrings[0], NULL);
        DeregisterEventSource(hEventSource);
    }
}


int execCmdUseGetline(const char* cmd, std::wofstream& ofs)
{
    FILE* pipe = _popen(cmd, "r");
    if (!pipe)
    {
        return -1;
    }
    HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    wchar_t buf[512] = { 0 };
    ofs << GetStrNow() << "\n";
    ofs.flush();
    while (fgetws(buf, 512, pipe) != NULL)
    {
        if (snapshot == INVALID_HANDLE_VALUE) {
            ofs << buf;
            ofs.flush();
        }
        else {
            std::wstring str(buf);
            int frel = str.find(L"TCP");
            int fre2 = str.find(L"UDP");
            if (frel != -1)
            {
                std::wistringstream iss(str);
                std::wstring protocol, localAddress, remoteAddress, state, processName;
                int processId;
                iss >> protocol >> localAddress >> remoteAddress >> state >> processId;

                PROCESSENTRY32 processEntry;
                processEntry.dwSize = sizeof(PROCESSENTRY32);
                // 遍历进程列表
                if (Process32First(snapshot, &processEntry)) {
                    do {
                        // 如果找到目标进程，则返回其名称
                        if (processEntry.th32ProcessID == processId) {
                            processName = processEntry.szExeFile;
                            break;
                        }
                    } while (Process32Next(snapshot, &processEntry));
                }
                ofs << protocol << "," << localAddress << "," << remoteAddress << "," << state << "," << processId << "," << processName << "\n";
                ofs.flush();
            }
            if (fre2 != -1) {
                std::wistringstream iss(str);
                std::wstring protocol, localAddress, remoteAddress, processName;
                int processId;
                // 读取并输出每个单词
                iss >> protocol >> localAddress >> remoteAddress >> processId >> processName;
                PROCESSENTRY32 processEntry;
                processEntry.dwSize = sizeof(PROCESSENTRY32);
                // 遍历进程列表
                if (Process32First(snapshot, &processEntry)) {
                    do {
                        // 如果找到目标进程，则返回其名称
                        if (processEntry.th32ProcessID == processId) {
                            processName = processEntry.szExeFile;
                            break;
                        }
                    } while (Process32Next(snapshot, &processEntry));
                }
                ofs << protocol << "," << localAddress << "," << remoteAddress << "," << "none" << "," << processId << "," << processName << "\n";
                ofs.flush();
            }
        }
    }
    ofs << "\n";
    ofs.flush();
    _pclose(pipe);
    if (snapshot != INVALID_HANDLE_VALUE)CloseHandle(snapshot);
    return 0;
}

int ConsoleMain()
{
    LOG(DEBUG) << "SimpleService: Service is running..." ;
    Packet* pk = new Packet();
    LOG(DEBUG) << "packet new" << std::endl;
    pk->Init();
    pk->Run();
    LOG(DEBUG) << "packet end" << std::endl;
    delete pk;
    pk = nullptr;

    return 0;
}

VOID WINAPI ServiceMain(DWORD argc, LPWSTR* argv)
{
    LOG(DEBUG) << "ServiceMain START...";
    g_ServiceStatus.dwCurrentState = SERVICE_START_PENDING;
    g_ServiceStatus.dwControlsAccepted = SERVICE_ACCEPT_STOP;
    // 注册服务控制处理函数
    g_ServiceStatusHandle = RegisterServiceCtrlHandler(SERVICE_NAME, ServiceCtrlHandler);
    if (g_ServiceStatusHandle == NULL)
    {
        LOG(ERROR) << "Handler not install";
        return;
    }

    SetServiceStatus(g_ServiceStatusHandle, &g_ServiceStatus);

    g_ServiceStatus.dwWin32ExitCode = S_OK;
    g_ServiceStatus.dwCheckPoint = 0;
    g_ServiceStatus.dwWaitHint = 0;
    g_ServiceStatus.dwCurrentState = SERVICE_RUNNING;
    SetServiceStatus(g_ServiceStatusHandle, &g_ServiceStatus);

    // 在 ServiceMain 中调用控制台程序的主函数
    ConsoleMain();

    // 设置服务状态为停止
    g_ServiceStatus.dwCurrentState = SERVICE_STOPPED;
    g_ServiceStatus.dwCheckPoint = 0;
    g_ServiceStatus.dwWaitHint = 0;

    if (!SetServiceStatus(g_ServiceStatusHandle, &g_ServiceStatus)){
        LOG(ERROR) << "SetServiceStatus failed, error code: " << GetLastError();
    }
    return;
}

void Init()
{
    g_ServiceStatusHandle = NULL;
    g_ServiceStatus.dwServiceType = SERVICE_WIN32_OWN_PROCESS | SERVICE_INTERACTIVE_PROCESS;
    g_ServiceStatus.dwCurrentState = SERVICE_START_PENDING;
    g_ServiceStatus.dwControlsAccepted = SERVICE_ACCEPT_STOP;
    g_ServiceStatus.dwWin32ExitCode = 0;
    g_ServiceStatus.dwServiceSpecificExitCode = 0;
    g_ServiceStatus.dwCheckPoint = 0;
    g_ServiceStatus.dwWaitHint = 0;
}


bool IsInstalled()
{
    SC_HANDLE hSCM = ::OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if (hSCM == NULL) {
        LOG(ERROR) << "OpenSCManager error..."<<GetLastError();
        return false;
    }

    SC_HANDLE hService = ::OpenService(hSCM, SERVICE_NAME, SERVICE_QUERY_CONFIG);
    if (hService == NULL) {
        ::CloseServiceHandle(hSCM);
        LOG(ERROR) << "OpenService error..."<<GetLastError() ;
        return false;
    }

    ::CloseServiceHandle(hService);
    ::CloseServiceHandle(hSCM);
    return false;
}


bool Install()
{
    if (IsInstalled())return true;
    SC_HANDLE hSCM = ::OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if (hSCM == NULL)return false;
    TCHAR szFilePath[MAX_PATH];
    ::GetModuleFileName(NULL, szFilePath, MAX_PATH);

    SC_HANDLE hService = ::CreateService(hSCM, SERVICE_NAME, SERVICE_NAME, SERVICE_ALL_ACCESS,
        SERVICE_WIN32_OWN_PROCESS | SERVICE_INTERACTIVE_PROCESS, SERVICE_AUTO_START, SERVICE_ERROR_NORMAL,
        szFilePath, NULL, NULL, _T(""), NULL, NULL);
    if (hService == NULL) {
        ::CloseServiceHandle(hSCM);
        LOG(ERROR) << "CreateService false...";
        return false;
    }

    // 创建服务后立即启动
    if (!StartService(hService, 0, NULL)){
        LOG(ERROR)  << " StartService false...";
    }
    else {
        printf("服务启动成功...");
        LOG(DEBUG) << " 服务启动成功...";
    }

    // 设置延迟自动启动信息
    SERVICE_DELAYED_AUTO_START_INFO delayedStartInfo{};
    delayedStartInfo.fDelayedAutostart = TRUE;

    // 修改服务配置  设置为延迟启动（开机1-2分钟内启动，否则自动启动会报错）
    if (!ChangeServiceConfig2(hService, SERVICE_CONFIG_DELAYED_AUTO_START_INFO, &delayedStartInfo)) {
        LOG(ERROR) << "Failed to change service configuration";
    }

    ::CloseServiceHandle(hService);
    ::CloseServiceHandle(hSCM);
    LOG(DEBUG) << "Install end...";
    return true;
}


bool Uninstall()
{
    SC_HANDLE schSCManager;
    SC_HANDLE schService;

    schSCManager = OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if (schSCManager == NULL){
        std::cout << "打开SCManager失败:" << GetLastError() << std::endl;
        return false;
    }

    schService = OpenService(schSCManager, TEXT("MyPcap"), SERVICE_ALL_ACCESS);
    if (schService == NULL) {
        std::cout << "打开服务失败：" << GetLastError() << std::endl;
        CloseServiceHandle(schSCManager);
        return false;
    }

    bool result = true;
    SERVICE_STATUS status;
    ::ControlService(schService, SERVICE_CONTROL_STOP, &status);
    if (!DeleteService(schService)){
        std::cout << "删除服务失败：" << GetLastError() << std::endl;
        result = false;
    }
    else{
        printf("服务删除成功");
    }

    CloseServiceHandle(schService);
    CloseServiceHandle(schSCManager);
    return result;
}



VOID WINAPI ServiceCtrlHandler(DWORD CtrlCode)
{
    switch (CtrlCode)
    {
    case SERVICE_CONTROL_STOP:
        //设置服务状态为停止
        g_ServiceStatus.dwCurrentState = SERVICE_STOP_PENDING;
        g_ServiceStatus.dwCheckPoint = 1;

        if (!SetServiceStatus(g_ServiceStatusHandle, &g_ServiceStatus))
        {
            std::cerr << "SetServiceStatus failed, error code: " << GetLastError() << std::endl;
        }
        Sleep(500);
        // 设置服务状态为停止
        g_ServiceStatus.dwCurrentState = SERVICE_STOPPED;
        g_ServiceStatus.dwCheckPoint = 0;
        g_ServiceStatus.dwWaitHint = 0;

        if (!SetServiceStatus(g_ServiceStatusHandle, &g_ServiceStatus))
        {
            std::cerr << "SetServiceStatus failed, error code: " << GetLastError() << std::endl;
        }

        PostThreadMessage(g_ThreadID, WM_CLOSE, 0, 0);
        break;
    case SERVICE_CONTROL_SHUTDOWN:
        exit(0);
        break;
    default:
        break;
    }

    return;
}


int APIENTRY _tWinMain(HINSTANCE hinstance, HINSTANCE hPrevInstance, LPTSTR lpCmdLine, int nCmdShow){
    EasyLogConf();    //easylog配置

    Init();
    g_ThreadID = ::GetCurrentThreadId();
    SERVICE_TABLE_ENTRY ServiceTable[] =
    {
        { SERVICE_NAME, (LPSERVICE_MAIN_FUNCTIONW)ServiceMain },
        { NULL, NULL }
    };

    if (_wcsicmp(lpCmdLine, _T("/install")) == 0)
    {
        LOG(DEBUG) << "install start...";
        Install();
    }
    else if (_wcsicmp(lpCmdLine, _T("/uninstall")) == 0)
    {
        LOG(DEBUG) << "uninstall start...";
        if (!Uninstall()) {
            LOG(DEBUG) << "uninstall error(" << GetLastError() << ")...";
        }
    }
    else {
        LOG(DEBUG) << "start...";
        if (!::StartServiceCtrlDispatcher(ServiceTable))
        {
            LOG(DEBUG) << "start error(" << GetLastError() << ")...";
        }
    }

    return 0;

}