/*************************************************************************
  > File Name: migrate.cpp
  > Author: TianLun Song
  > Mail: songtianlun@frytea.com
  > Blog: https://blog.frytea.com
  > Created Time: Wed 23 Dec 2020 07:10:09 PM CST
 ************************************************************************/

#include <criu/criu.h>
#include<iostream>
#include<sys/types.h>
#include<sys/stat.h>
#include<fcntl.h>
#include<unistd.h>
#include<string.h>
#include<stdio.h>
#include<stdlib.h>

using namespace std;

class CMigrateApp
{
    public:
        CMigrateApp(int iPid = 0);
        CMigrateApp(char *sPid);
        ~CMigrateApp();
        int Checkout();
        int Restore();
        int GetPid();

    private:
        int GetPidByCMD(char * sCMD);

    private:
        char criu_dir[64];
        char criu_server[64];
        char criu_log[64];
        char criu_cmd[64];
        int  criu_log_level;
        int  criu_pid;
};


CMigrateApp::CMigrateApp(int iPid)
{
    strcpy(criu_dir, "./migrate_imgs");
    strcpy(criu_server, "/root/criu_service.socket");
    strcpy(criu_log, "criu.log");
    criu_log_level = 4;
    criu_pid = iPid;

    

    criu_init_opts();
    criu_set_service_address(criu_server);
    int fd = open(criu_dir, O_DIRECTORY);
    criu_set_images_dir_fd(fd);
    
    criu_set_log_file(criu_log);
    criu_set_log_level(criu_log_level);
    
    criu_set_shell_job(true);
}

CMigrateApp::CMigrateApp(char* sPid)
{
    char sMkdir[64] = {0};

    strcpy(criu_dir, "/root/migrate_imgs/");
    strcpy(criu_server, "/root/criu_service.socket");
    strcpy(criu_log, "criu.log");
    criu_log_level = 4;

    strcpy(criu_cmd, sPid);
    strcat(criu_dir, sPid);
    sprintf(sMkdir,"%s %s", "mkdir -p", criu_dir);
    system(sMkdir);

    criu_pid = GetPidByCMD(criu_cmd);
    //cout << "Pid: " << criu_pid << endl;

    
    criu_init_opts();
    criu_set_service_address(criu_server);
    int fd = open(criu_dir, O_DIRECTORY);
    criu_set_images_dir_fd(fd);
    
    criu_set_log_file(criu_log);
    criu_set_log_level(criu_log_level);
    
    criu_set_shell_job(true);
    criu_set_ext_unix_sk(true);
}

int CMigrateApp::Checkout()
{
    if(0 >= criu_pid)
    {
        //cout << "Error Pid! Dump failed!" << endl;
        return -1;
    }
    criu_set_pid(criu_pid);
    return criu_dump();
}

int CMigrateApp::Restore()
{
    return criu_restore();
}

int CMigrateApp::GetPidByCMD(char *sCMD)
{
    char cPid[62] = {0}; 
    char sSeekCmd[64] = {0};
    FILE *fPid = NULL; 
    int iPid = -1;
    sprintf(sSeekCmd, "%s %s %s", "pgrep", sCMD, "| head -1");
    fPid = popen(sSeekCmd, "r");
    if(NULL == fPid)
        return -1;
    fgets(cPid, sizeof(cPid), fPid);
    iPid = atoi(cPid);
    //cout << "SeekCmd: " << sSeekCmd << ", Pid: " << iPid << ", cPid: " << cPid << endl;
    return iPid;
}
int CMigrateApp::GetPid()
{
    return criu_pid;
}

int main()
{
    char piggie_cmd[21][64] = {
        "arpd","dhcp","lldpd","ntpclient","ripd","snmpd","udpsvd",
        "brctl","l2monitor","master ","ospfd","ripngd","snmptrap","vtysh",
        "configmanage","lag","mstp","pidmonitor","rmon","udhcpd","zebra"};
    CMigrateApp * migrate_app[21];

    for(int i = 0; i < 21; i++)
    {
        migrate_app[i] = new CMigrateApp(piggie_cmd[i]);
        // int iRet = 0;
        cout << "Dump " << piggie_cmd[i] << "[ID:" << migrate_app[i]->GetPid()  << "]" << " Result: " << migrate_app[i]->Checkout() << endl;
        //sleep(1);
        //cout << "Restore " << piggie_cmd[i] << " Result: " << migrate_app[i]->Restore() << endl;
        //sleep(5);
    }
    return 0;
}