#ifndef WIN32
#include <unistd.h>
#include <sys/socket.h>
#include <algorithm>
#include "shared.h"
#include "proc_mon_linux.h"
#include <malloc.h>
#include "linux_proc_baseline.h"


#define MAX_PIPE_WAIT_USECOND 100
#define PROCESS_SCAN_ITERTIME  10
const char *WM_PROC_LOCATION = "proc_collector";

ProcMon* ProcMon::m_instance = NULL;

ProcMon& ProcMon::instance()
{
    if ( m_instance == NULL )
    {
        m_instance = new ProcMon;
    }

    return *m_instance;
}

void ProcMon::destroy()
{
    delete m_instance;
    m_instance = NULL;
}

int ProcMon::init( const wm_proc_t* config )
{
    pthread_mutex_init(&m_mutex, NULL);
    setConfig( config );

    int ret = setupNetlink();
    if (ret != 0 )
    {
        fini();
        return -1;
    }

    ret = setupMQ();
    if (ret != 0 )
    {
        fini();
        return -1;
    }

    return 0;
}

int ProcMon::working()
{
    int rc;
    event_msg_t procMsg;
    fd_set readfds;
    int maxfd = m_procFd + 1;
    struct timeval tv;
    time_t last = time(0);

    if ( m_config.flags.enabled == 0 )
    {
        return 0;
    }

    getInventoryProcesses();


    while ( 1 )
    {
        FD_ZERO(&readfds);
        FD_SET( m_procFd, &readfds );
        tv.tv_sec = 1;
        tv.tv_usec = 0;

        if ( time(0) - last > PROCESS_SCAN_ITERTIME )
        {
            // scan processes
            scanProcesses(last);
            last = time(0);
        }

        rc = select( maxfd, &readfds, NULL, NULL, &tv );
        if ( 0 == rc )
        {
            continue;
        }

        if ( -1 == rc )
        {
            if ( errno == EINTR )
            {
                continue;
            }

            mterror(WM_PROC_LOGTAG, "failed to listen to netlink socket, errno=(%d:%m)", errno);
            return rc;
        }

        if ( FD_ISSET( m_procFd, &readfds ) )
        {
            rc = recv( m_procFd, &procMsg, sizeof(procMsg), 0 );
            if( rc > 0 )
            {
                switch (procMsg.proc_ev.what)
                {
                case proc_event::PROC_EVENT_FORK:
                    handleForkMsg(procMsg);
                    break;
                case proc_event::PROC_EVENT_EXEC:
                    handleExecMsg(procMsg);
                    break;
                case proc_event::PROC_EVENT_COMM:
                    break;
                case proc_event::PROC_EVENT_EXIT:
                    handleExitMsg(procMsg);
                    break;
                default:
                    break;
                }
            }
            else if (rc == -1)
            {
                if (errno == EINTR)
                {
                    continue;
                }
                mterror(WM_PROC_LOGTAG, "failed to received from netlink socket, errno=(%d:%m)", errno);
            }
        }

		malloc_trim(0);
    }
}

ProcMon::ProcMon() 
{
    m_mutex = PTHREAD_MUTEX_INITIALIZER;
    m_processes.clear();
    m_pidCache.clear();
    m_procFd = -1;
    m_queueFd = -1;
}

ProcMon::~ProcMon()
{
    fini();
}

void ProcMon::setConfig( const wm_proc_t* config  )
{
    m_config.flags.enabled = config->flags.enabled;
    m_config.flags.rootcheck = config->flags.rootcheck;
    m_config.flags.limitcheck = config->flags.limitcheck;
    m_config.flags.environcheck = config->flags.environcheck;
    m_config.flags.statuscheck = config->flags.statuscheck;
    m_config.flags.threadcheck = config->flags.threadcheck;
    m_config.flags.cgroupcheck = config->flags.cgroupcheck;
    m_config.flags.libcheck = config->flags.libcheck;
    m_config.flags.malwarecheck = config->flags.malwarecheck;
    m_config.interval = config->interval;
}

int ProcMon::setupNetlink()
{
    int rc;
    int nl_sock;
    struct sockaddr_nl sa_nl;
    register_msg_t nlcn_msg;

    nl_sock = socket(PF_NETLINK, SOCK_DGRAM, NETLINK_CONNECTOR);
    if (nl_sock == -1)
    {
        mterror(WM_PROC_LOGTAG, "Can't open netlink socket");
        return -1;
    }

    sa_nl.nl_family = AF_NETLINK;
    sa_nl.nl_groups = CN_IDX_PROC;
    sa_nl.nl_pid = getpid();

    rc = bind(nl_sock, (struct sockaddr *)&sa_nl, sizeof(sa_nl));
    if (rc == -1)
    {
        mterror(WM_PROC_LOGTAG, "Can't bind netlink socket");
        close(nl_sock);
        return -1;
    }

    // create listener
    memset(&nlcn_msg, 0, sizeof(nlcn_msg));
    nlcn_msg.nl_hdr.nlmsg_len = sizeof(nlcn_msg);
    nlcn_msg.nl_hdr.nlmsg_pid = getpid();
    nlcn_msg.nl_hdr.nlmsg_type = NLMSG_DONE;

    nlcn_msg.cn_msg.id.idx = CN_IDX_PROC;
    nlcn_msg.cn_msg.id.val = CN_VAL_PROC;
    nlcn_msg.cn_msg.len = sizeof(enum proc_cn_mcast_op);

    nlcn_msg.cn_mcast = PROC_CN_MCAST_LISTEN;

    rc = send(nl_sock, &nlcn_msg, sizeof(nlcn_msg), 0);
    if (rc == -1)
    {
        mterror(WM_PROC_LOGTAG, "can't register to netlink");
        close(nl_sock);
        return -1;
    }

    m_procFd = nl_sock;
    return 0;
}

int ProcMon::setupMQ()
{
    unsigned int indx;
    // Connect to socket
    for (indx = 0; indx < WM_MAX_ATTEMPTS; indx++)
    {
        m_queueFd = StartMQ(DEFAULTQPATH, WRITE);
        if (m_queueFd > 0)
        {
            break;
        }
        wm_delay(1000 * WM_MAX_WAIT);
    }

    if (indx == WM_MAX_ATTEMPTS)
    {
        mterror(WM_PROC_LOGTAG, "Can't connect to queue.");
        pthread_exit(NULL);
    }

    return 0;
}


void ProcMon::fini()
{
    m_processes.clear();
    m_pidCache.clear();
    if (m_procFd != -1 )
    {
        close( m_procFd);
        m_procFd = -1;
    }

    if (m_queueFd != -1 )
    {
        close( m_queueFd);
        m_queueFd = -1;
    }
}

void ProcMon::getInventoryProcesses()
{
    DIR *procDir = opendir( "/proc" );
    if ( NULL == procDir )
    {
        mterror(WM_PROC_LOGTAG, "failed to open /proc, errno=(%d:%m)", errno);
        return;
    }

    struct dirent *entry = NULL;
    while (NULL != (entry = readdir(procDir)))
    {
        if (!strcmp(entry->d_name, ".") || !strcmp(entry->d_name, ".."))
        {
            continue;
        }
        int pid = atoi(entry->d_name);
        if (0 == pid)
        {
            continue;
        }

        Process *proc = new Process(pid);
        if ( NULL == proc )
        {
            continue;
        }
	    if(0 != proc->init())
	    { 
	        delete proc;
			proc = NULL;
			continue;
	    }
		std::pair<std::map<pid_t, Process*>::iterator,bool> ret = m_processes.insert(std::make_pair(pid, proc));
		if(!ret.second)
		{	
			delete proc;
			proc = NULL;
			continue;
		}
        std::vector<std::string> events = proc->getEvents();
        if ( !events.empty() )
        {
            reportEvents(events);
        }
        usleep( 500 );
    }

    return;
}

void ProcMon::handleForkMsg( const event_msg_t& procMsg )
{
    pid_t pid = procMsg.proc_ev.event_data.fork.child_pid;
    auto iter = m_pidCache.find( pid );
    if ( iter != m_pidCache.end() )
    {
        m_pidCache.erase(iter);
        return;
    }

    PidRelation relation;
    relation.pid = pid;
    relation.ppid = procMsg.proc_ev.event_data.fork.parent_pid;
    relation.pidEvent = PidRelation::FORKEVENT;
    relation.born = time(0);

    m_pidCache.insert( std::make_pair(relation.pid, relation ));
}

void ProcMon::handleExecMsg( const event_msg_t& procMsg )
{
    pid_t pid = procMsg.proc_ev.event_data.exec.process_pid;
    auto iter = m_pidCache.find( pid );
    if ( iter == m_pidCache.end() )
    {
        return;
    }

    //mterror(WM_PROC_LOGTAG, "cache size:%lu", m_pidCache.size());
    Process *proc = new Process( pid, iter->second.ppid, iter->second.born );
    if ( NULL == proc )
    {
        return;
    }
    // check hidden process 
    if ( 0 != proc->init()  )
    {
        return;
    }
    m_processes.insert( std::make_pair(pid, proc ));
    std::vector<std::string> events = proc->getEvents();
    if (!events.empty())
    {
        reportEvents(events);
    }
//	m_pidCache.erase(iter);
    //mterror(WM_PROC_LOGTAG, "exec process:%d exe:%s", pid, proc.exe().c_str());
}

void ProcMon::handleExitMsg( const event_msg_t& procMsg )
{
    pid_t pid = procMsg.proc_ev.event_data.exit.process_pid;
    auto iter = m_pidCache.find( pid );
    if ( iter != m_pidCache.end() )
    {
        m_pidCache.erase(iter);
    }

    auto procIter = m_processes.find(pid);
    if ( procIter == m_processes.end() )
    {
        return;
    }

    std::vector<int> signals{ 3, 9, 10, 12, 15,16,17,30,31};

    //if ( signals.end() == std::find( signals.begin(), signals.end(), procMsg.proc_ev.event_data.exit.exit_signal )
    //|| (procMsg.proc_ev.event_data.exit.exit_code != 0))
    {
        // report abnormal events
        cJSON *obj= cJSON_CreateObject();
        cJSON_AddStringToObject(obj, "type", "exit");
        cJSON_AddStringToObject(obj, "exe", procIter->second->exe().c_str());
        cJSON_AddNumberToObject(obj, "pid", procIter->first);
        cJSON_AddNumberToObject(obj, "event_time", time(0));
        cJSON_AddNumberToObject(obj, "exit_signal", procMsg.proc_ev.event_data.exit.exit_signal);
        cJSON_AddNumberToObject(obj, "exit_code", procMsg.proc_ev.event_data.exit.exit_code);
        char *data = cJSON_PrintUnformatted( obj );
        std::string msg(data);
        free(data);
        reportEvent( msg );
        cJSON_Delete(obj);
        //mterror(WM_PROC_LOGTAG, "%s", msg.c_str());
    }

    delete procIter->second;
    m_processes.erase(procIter);
    //mterror(WM_PROC_LOGTAG, "process collection:%d", m_processes.size());
}

void ProcMon::scanProcesses(time_t last)
{
    time_t now = time(0);
//	mterror(WM_PROC_LOGTAG, "pid cache size: %d , process collection size : %d", m_pidCache.size(), m_processes.size());
    for (auto iter = m_pidCache.begin(); iter != m_pidCache.end(); )
    {
        if ( now - iter->second.born < 1 || m_processes.find(iter->second.pid) != m_processes.end())
        {
            iter++;
        }
        else
        {
            Process *proc = new Process( iter->second.pid, iter->second.ppid, iter->second.born);
            if (NULL != proc ) 
            {
            	if(0 != proc->init())
            	{ 
            		delete proc;
					proc = NULL;
					continue;
            	}
                std::pair<std::map<pid_t, Process*>::iterator, bool> ret = m_processes.insert(std::make_pair(iter->first, proc));
				if(!ret.second)
				{	
					delete proc;
					proc = NULL;
					continue;
				}
                std::vector<std::string> events = proc->getEvents();
                if (!events.empty())
                {
                    reportEvents(events);
                }

            }
            else
            {
                 // hidden process found
                 cJSON *obj = cJSON_CreateObject();
                 cJSON_AddStringToObject(obj, "type", "rootkit");
                 cJSON_AddNumberToObject(obj, "pid", iter->first);
                 cJSON_AddNumberToObject(obj, "event_time", iter->second.born);
                 cJSON_AddNumberToObject(obj, "ppid", iter->second.ppid);
                 auto parent = m_processes.find(iter->second.ppid);
                 if (parent != m_processes.end())
                 {
                     cJSON_AddStringToObject(obj, "parent exe", parent->second->exe().c_str());
                 }
                 char *data = cJSON_PrintUnformatted(obj);
                 std::string msg(data);
                 free(data);
                 reportEvent(msg);
                 cJSON_Delete(obj);
            }
            iter = m_pidCache.erase(iter);
        }
        usleep( 500 );
    }
/*
    if ( now - last < 60 )
    {
        return;
    }
*/
    for (auto& procIter : m_processes )
    {
        procIter.second->scan();
        std::vector<std::string> events = procIter.second->getEvents();
        if (!events.empty())
        {
            reportEvents(events);
        }
        usleep( 5000 );
    }
}

void ProcMon::reportEvent( const std::string& event )
{
    wm_sendmsg(MAX_PIPE_WAIT_USECOND, m_queueFd, event.c_str(), WM_PROC_LOCATION, PROC_COLLECTOR_MQ);
}

void ProcMon::reportEvents(const std::vector<std::string> &events)
{
    for (auto it : events)
    {
        wm_sendmsg(MAX_PIPE_WAIT_USECOND, m_queueFd, it.c_str(), WM_PROC_LOCATION, PROC_COLLECTOR_MQ);
    }
}
#endif