malware_handler.cpp

#include <iostream>
#include <cstring>
#include <unistd.h>
#include <map>
#include <vector>
#include <cstdio>
using namespace std;

#define THRESHOLD 50 
//get total cpu time of all processes in system in seconds
long double get_total_usage(){
	FILE *fp = fopen("/proc/stat","r");
	if(!fp){
		perror("fopen");
		exit(0);
	}
	char buff[2048];
	fgets(buff,sizeof(buff),fp);
	char *p = strtok(buff," ");
	int i = 0;
	long double ans=0;
	while(p){
		if(i>0)ans += atoi(p);
		p = strtok(NULL," ");
		i++;
	}
	return ans/sysconf(_SC_CLK_TCK);
}

//get usage of process <pid> in seconds
long double get_usage(int pid){
	char file_name[100];
    snprintf(file_name,sizeof(file_name), "/proc/%d/stat", pid);
    FILE *fp = fopen(file_name, "r");
    if (!fp){
        perror("proc file");
        exit(1);
    }

    char buf[2048];
    fgets(buf, sizeof(buf), fp);
    char *p = strtok(buf, " ");
    long long i=0,u_time, s_time;
    while(p){
	    if(i == 13)u_time=strtol(p, NULL, 10);
        else if(i == 14)s_time=strtol(p, NULL, 10);
        p = strtok(NULL, " ");
        i++;
    }
	fclose(fp);
	return (1.0*(u_time + s_time))/sysconf(_SC_CLK_TCK);
}

void usage_stats(int pid,int suggest_flag,map <long long,long double> &mp,vector<long long> &parents,bool flag){
	int i=0;
	char path[256];
	mp[pid] = get_usage(pid);
	if(flag)cout << "Current process id " << pid << endl;
	parents.clear();
	parents.push_back(pid);	
	while(1){
	  	snprintf(path, sizeof(path), "/proc/%d/status", pid);
		FILE* file = fopen(path, "r");
		if(!file){
			perror("fopen");
			return;
		}
		char line[256];
  		while(fgets(line, sizeof(line), file)){
    		if(strncmp(line, "PPid:", 5) == 0){
      			int ppid;
      			sscanf(line + 5, "%d", &ppid); 
				if(flag)cout << "Parent process ID: " << ppid << endl;
      			parents.push_back(ppid);
				if(ppid==0){
					pid=ppid;
					break;
				}
				mp[ppid] = get_usage(ppid);

				pid = ppid;
				break;
    		}
  		}
		if(pid==0)break;
		fclose(file);
		i++;
	}	
	long double total_time = get_total_usage();
	mp[0] = total_time;			
}

int main(int argc,char *argv[]){
	int pid,suggest_flag=0;
	if(argc==2){
		pid = atoi(argv[1]);
	}
	else if(argc==3){
		pid = (strcmp(argv[1],"-suggest")==0)? 2:1;
		suggest_flag=1;
		pid = atoi(argv[pid]);
	}
	map <long long,long double> mp1,mp2,process_time;
	mp1.clear();
	mp2.clear();
	vector <long long> v1;
	usage_stats(pid,suggest_flag,mp1,v1,1);
	if(suggest_flag==0){
		exit(0);
	}
	sleep(5);
	usage_stats(pid,suggest_flag,mp2,v1,0);
	long double total_usage = mp2[0] - mp1[0];
	for(auto &it:mp1){
		if(it.first == 0)continue;
		if(mp2.find(it.first)==mp2.end()){
			continue;
		}
		// cout << "Process " << it.first << " CPU % " << ((mp2[it.first]-it.second)/total_usage) * 100.0 << endl;	
		process_time[it.first] = ((mp2[it.first]-it.second)/total_usage) * 100.0;
	}

	int len = v1.size();
	if(len <= 2){
		printf("Please check pid again using htop!\n");
		exit(0);
	}	
	long double curr_usage,prev_usage = -1;
	for(int i=0;i<len;i++){
		if(i+1==len)break;
		curr_usage = process_time[v1[i]];
		if(prev_usage<0){
			prev_usage = curr_usage;
			continue;
		}
		
		if(curr_usage < (1.0*prev_usage)/THRESHOLD){
			cout << "Process ID for the parent malware process P is: " << v1[i] << endl;
			exit(0);
		}
		prev_usage = curr_usage;
	}		
	cout << "Please re-check pid using htop, no suitable processes found" << endl;
	return 0; 
}