/*************************************************************************
  > 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<iostream>
#include<algorithm>  /* sort() */
#include<sys/types.h>
#include<sys/stat.h>
#include<fcntl.h>
#include<unistd.h>
#include<string.h>
#include<stdio.h>
#include<stdlib.h>
#include<sys/types.h>
#include<sys/socket.h>

class CMigrateApp
{
#define CRIU_BINARY         "/usr/sbin/criu"
#define CRIU_DIR            "/tmp/criu"
#define IMG_DIR	        	"images"
#define DUMP_LOG_FILE   	"dump.log"
#define RESTORE_LOG_FILE    "restore.log"
#define RESTORE_PID_FILE    "restore.pid"
#define INHERIT_FD_OPTION   "--inherit-fd"
#define OLD_LOG_FILE        "/tmp/criu/oldlog"
#define NEW_LOG_FILE        "/tmp/criu/newlog"

public:
    // CMigrateApp(int iPid = 0);
    CMigrateApp(char *sCMD, int iPid = 0);
    ~CMigrateApp(){}
    char *GetCMD();
    char *GetImgDir();
    int  Checkout();
    int  Restore();
    int  GetPid();
    int  GetUnixNum();

private:
    int  GetPidByCMD(char * sCMD);
    int  SetPidByCMD(char * sCMD);
    void GetCMDByPid(int iPid, char *sCMD);
    int  isUnixActivity(unsigned int iNode, char * cPid);
    int  SetUnixByPid(int iPid);
    void Mkdir_safe(char *dirname, int mode);

private:
    int  i_criu_pid;
    int  i_unix_num;
    char s_criu_cmd[64];
    char s_criu_pid[8];
    char s_criu_img[64];
    char s_criu_dump_log[64];
    char s_criu_restore_log[64];
    char s_criu_restore_pidfile[64];
    char s_external_sk_ino[1024];
    char s_inh_unixsk_arg[1024];
    const char *s_dump_argv[15] = {
        "criu", "dump", "-j",
        "-D", s_criu_img, "-o", s_criu_dump_log,
        "-v4", "-x", "--lib /root/stl/lib",
        "--tcp-established",
        s_external_sk_ino,
        "-t", s_criu_pid,
        NULL };
    const char *s_restore_argv[18] = {
        "criu", "restore", "-d","-j",
        "-D", s_criu_img, "-o", s_criu_restore_log,
        "--pidfile", s_criu_restore_pidfile,
        "-v4", "-x", "--lib /root/stl/lib",
        "--tcp-established",
        s_inh_unixsk_arg,
        NULL };
};

CMigrateApp::CMigrateApp(char* sCMD, int iPid)
{
    char sMkdir[64] = {0};

    // init data
    memset(s_criu_img, 0x00, sizeof s_criu_img);
    memset(s_criu_dump_log, 0x00, sizeof s_criu_dump_log);
    memset(s_criu_restore_log, 0x00, sizeof s_criu_restore_log);
    memset(s_criu_restore_pidfile, 0x00, sizeof s_criu_restore_pidfile);
    memset(s_criu_cmd, 0x00, sizeof s_criu_cmd);
    i_criu_pid = 0;
    i_unix_num = 0;
    memset(s_criu_pid, 0x00, sizeof s_criu_pid);
    memset(s_external_sk_ino, 0x00, sizeof s_external_sk_ino);
    memset(s_inh_unixsk_arg, 0x00, sizeof s_inh_unixsk_arg);

    if( 0 != iPid )
    {
        GetCMDByPid(iPid, s_criu_cmd);
        i_criu_pid = iPid;
        snprintf(s_criu_pid, sizeof s_criu_pid, "%d", i_criu_pid);
    }
    else if( NULL != sCMD )
    {
        strcpy(s_criu_cmd, sCMD);
        SetPidByCMD(s_criu_cmd);
    }
    else
    {
        printf("Error of [iPid: %d] or [sCMD: %s] ", iPid, sCMD);
    }

    snprintf(s_criu_img, sizeof s_criu_img, "%s/%s/%s", CRIU_DIR, s_criu_cmd, IMG_DIR);
    snprintf(s_criu_dump_log, sizeof s_criu_dump_log, "%s/%s/%s", CRIU_DIR, s_criu_cmd, DUMP_LOG_FILE);
    snprintf(s_criu_restore_log, sizeof s_criu_restore_log, "%s/%s/%s", CRIU_DIR, s_criu_cmd, RESTORE_LOG_FILE);
    snprintf(s_criu_restore_pidfile, sizeof s_criu_restore_pidfile, "%s/%s/%s", CRIU_DIR, s_criu_cmd, RESTORE_PID_FILE);
    i_unix_num = SetUnixByPid(i_criu_pid);
    // printf("PCMD: %s, Unix Sum: %d\n",sCMD, i_unix_num);

    sprintf(sMkdir,"mkdir -p %s &", s_criu_img);
    system(sMkdir);
    printf("App will be to Migrate: [iPid: %d] [CMD: %s] [Unix Num: %d]\n", i_criu_pid, s_criu_cmd, i_unix_num);
}

int CMigrateApp::Checkout()
{
    char sCmd[2048] = {0};
    char sCpCMd[64] = {0};
    if(0 >= i_criu_pid)
    {
        printf("Error Pid! Dump failed!\n");
        return -1;
    }
    
    for (int i = 0; s_dump_argv[i] != NULL; i++)
        sprintf(sCmd, "%s %s", sCmd, s_dump_argv[i]);
    strcat(sCmd,"&");
    printf("Dump [CMD:%s] [PID:%d] [Unix Num:%d] cmd: %s\n",s_criu_cmd,i_criu_pid, i_unix_num, sCmd);
    system(sCmd);

    snprintf(sCpCMd, sizeof sCpCMd, "cp /tmp/1 /tmp/criu/%s/", s_criu_cmd);
    system(sCpCMd);
    return 0;
}

int CMigrateApp::Restore()
{
    char sCmd[2048] = {0};
    char sCpCmd[64] = {0};
    if(0 >= i_criu_pid)
    {
        printf("Error Pid! Restore neet to right dump!\n");
        return -1;
    }
    for (int i = 0; s_restore_argv[i] != NULL; i++)
        sprintf(sCmd, "%s %s", sCmd, s_restore_argv[i]);
    strcat(sCmd,"&");
    /* 进程有几个socket就创建几次 */
    for (int i = 0; i < i_unix_num; i++)
        socket(AF_UNIX, SOCK_STREAM, 0);
    /* 将该进程转储时的/tmp/1 复制到指定路径下 */
    snprintf(sCpCmd, sizeof sCpCmd, "cp /tmp/criu/%s/1 /tmp/1", s_criu_cmd);
    system(sCpCmd);
    system("chmod 644 /tmp/1");
    /* 恢复 */
    printf("Restore [CMD:%s] [Unix Num:%d] cmd: %s\n",s_criu_cmd, i_unix_num, sCmd);
    system(sCmd);
    return 0;
}

int CMigrateApp::GetPidByCMD(char *sCMD)
{
    char cPid[62] = {0}; 
    char sSeekCmd[64] = {0};
    FILE *fPid = NULL; 
    int iPid = -1;
    sprintf(sSeekCmd, "pgrep %s | head -1",sCMD);
    fPid = popen(sSeekCmd, "r");
    if(NULL == fPid)
        return -1;
    fgets(cPid, sizeof(cPid), fPid);
    iPid = atoi(cPid);
    // printf("SeekCmd: %s, Pid: %d, cPid: %s\n",sSeekCmd,iPid,cPid);
    return iPid;
}
void CMigrateApp::GetCMDByPid(int iPid, char *sCMD)
{
    FILE *fcmd = NULL; 
    char sGetCMD[64] = {0};
    snprintf(sGetCMD, sizeof sGetCMD, "grep 'Name' /proc/%d/status | awk '{print $2}'", iPid);
    fcmd = popen(sGetCMD, "r");
    fgets(sCMD, sizeof(sCMD), fcmd);
}
int CMigrateApp::SetPidByCMD(char *sCMD)
{
    char cPid[62] = {0}; 
    char sSeekCmd[64] = {0};
    FILE *fPid = NULL; 
    int iPid = -1;
    sprintf(sSeekCmd, "pgrep %s | head -1", sCMD);
    fPid = popen(sSeekCmd, "r");
    if(NULL == fPid)
        return -1;
    fgets(cPid, sizeof(cPid), fPid);
    iPid = atoi(cPid);
    
    i_criu_pid = atoi(cPid);
    snprintf(s_criu_pid, sizeof s_criu_pid, "%d", iPid);
    // strcpy(s_criu_pid, cPid);
    pclose(fPid);
    //printf("SeekCmd: %s, Pid: %d, cPid: %s\n",sSeekCmd,iPid,cPid);
    return iPid;
}
int CMigrateApp::SetUnixByPid(int iPid)
{
    char sUnixStatCmd[64] = {0};
    char sNowCriuUnixCmd[64] = {0};
    char sNosRestoreUnixCmd[64] = {0};
    FILE *fpUnixStat;
    int iFdNum = 3;
    char tmp[2048]; //存储每一行输出
    unsigned int iNode;
    int iSumUnix = 0;
    // int i = 0;

    if( iPid <=0 )
        return 0;

    sprintf(sUnixStatCmd, "cat /proc/%d/net/unix | awk '{print $7}'", iPid);
    fpUnixStat = popen(sUnixStatCmd, "r");
    if (fpUnixStat == NULL) {
        return -1;
    }
    // 跳过第一条无效数据
    fgets(tmp, sizeof(tmp) - 1, fpUnixStat);
    // i = 0;
    // 输出检查状态
    // printf("check unix ..");
    // fflush(stdout);
    while (fgets(tmp, sizeof(tmp) - 1, fpUnixStat)) {
        // 输出检查状态
        // switch(i++){
        //     case 1:
        //         printf(".");
        //         fflush(stdout);
        //         break;
        //     case 5:
        //         printf("\b");
        //         fflush(stdout);
        //         break;
        //     case 10:
        //         i=0;
        //     default:
        //         break;
        // };
        sscanf(tmp, "%u", &iNode);
        if(!isUnixActivity(iNode,s_criu_cmd))
            continue;
        if(iNode==0)
            continue;
        // printf("iPid: %d, Unix iNode: %d, is activity.\n", iPid,iNode);
        snprintf(sNowCriuUnixCmd, sizeof(sNowCriuUnixCmd), "--ext-unix-sk=%u ",
			(unsigned int)iNode);
        snprintf(sNowCriuUnixCmd, sizeof(sNowCriuUnixCmd), "--external unix[%u] ",
			(unsigned int)iNode);
        snprintf(sNosRestoreUnixCmd, sizeof sNosRestoreUnixCmd, "--inherit-fd fd[%d]:socket:[%u] ", 
            iFdNum++,(unsigned int)iNode);
        strcat(s_external_sk_ino, sNowCriuUnixCmd);
        strcat(s_inh_unixsk_arg, sNosRestoreUnixCmd);
        iSumUnix++;
    }
    // printf("\r"); // 清空 check unix 状态字符打印
    pclose(fpUnixStat);
    return iSumUnix;
}
int CMigrateApp::isUnixActivity(unsigned int iNode, char * cPid)
{
    char sCheckUnixCmd[64] = {0};
    FILE *fpCheckUnix;
    char tmp[2048];
    // 检查当前 unix socket 是否处于活跃状态
    sprintf(sCheckUnixCmd, "netstat -xap | grep %s | grep %u", cPid, iNode);
    fpCheckUnix = popen(sCheckUnixCmd, "r");
    if (fpCheckUnix == NULL) {
        return -1;
    }
    memset(tmp, 0x0, sizeof(tmp));
    fgets(tmp, sizeof(tmp) - 1, fpCheckUnix);
    pclose(fpCheckUnix);
    printf("Check Cmd: %s,  Rst len: %d", sCheckUnixCmd, (int)strlen(tmp));
    fflush(stdout);
    printf("%40s\r"," "); // 使用空格清理脏字符并回退光标至行首
    fflush(stdout);
    if(strlen(tmp)>1)
        return 1;
    return 0;
}
int CMigrateApp::GetPid()
{
    return i_criu_pid;
}
int CMigrateApp::GetUnixNum()
{
    return i_unix_num;
}
char *CMigrateApp::GetCMD()
{
    return s_criu_cmd;
}
char *CMigrateApp::GetImgDir()
{
    return s_criu_img;
}

void CMigrateApp::Mkdir_safe(char *dirname, int mode)
{
	if (mkdir(dirname, mode) == -1 && errno != EEXIST)
		printf("Success to mkdir dirname=%s mode=0x%x\n", dirname, mode);
    else
        printf("Fail to make dir %s", dirname);
}

// 自定义对象比较函数（根据unix个数排序）
bool cmp(CMigrateApp a, CMigrateApp b)
{
    return a.GetUnixNum() < b.GetUnixNum();
}
void MigrateMySelf()
{
    char sRestoreCmd[128] = {0};
    CMigrateApp CMigrateMySelf(NULL, getpid());
    snprintf(sRestoreCmd, sizeof sRestoreCmd, "criu restore -D %s -j", CMigrateMySelf.GetImgDir());
    printf("Run cmd when restore: %s\n", sRestoreCmd);
    CMigrateMySelf.Checkout();
}

int main()
{
    char piggie_cmd[21][64] = {
        /* create ipc unix socket */
        "arpd", "dhcp", "lldpd", "ospfd", "ripd", "snmptrap", "rmon", "ripngd", "mstp", "lag","master","zebra","pidmonitor",
        "ntpclient", "snmpd", "udpsvd", "brctl", "l2monitor", "vtysh", "configmanage", "udhcpd"};
    std::vector <CMigrateApp> vMigrateApps;

    // 清空进程检查点数据文件夹，避免遗留文件影响进程恢复
    system("rm -rf /tmp/criu/*");

    for(int i = 0; i < 21; i++)
    {
        CMigrateApp iC(piggie_cmd[i]);
        vMigrateApps.push_back(iC);
        // printf("%s [PID %d] will be push back, vector size: %d\n", iC.GetCMD(),iC.GetPid(), (int)vMigrateApps.size());
    }
    
    printf("-----------------------after sort-----------------------\n");
    for(auto v : vMigrateApps)
    {
        printf("[%s:%d] has %d unix socket\n", v.GetCMD(), v.GetPid(), v.GetUnixNum()); 
    }
    // sort(vMigrateApps.begin(), vMigrateApps.end(),cmp);
    printf("--------------------------  dump  -------------------------\n");
    for(auto v : vMigrateApps)
    {
        v.Checkout();
    }
    for(int j=10;j>0;j--)
    {
        printf("Waiting for start restore....%d\n",j);
        if(j==5)
            MigrateMySelf();
        sleep(1);
    }
    
    printf("Enter any character to continue\n"); 
    getchar();
    printf("------------------------  restore  ------------------------\n");
    for(auto v : vMigrateApps)
    {
        if(v.GetPid()<=0)
        {
            printf("failed to dump, jump %s.\n", v.GetCMD());
            continue;
        }
        v.Restore();
    }
    return 0;
}