/* KMEnter.cpp
 *
 * Copyright (C) 2024 OpenKylin, Inc
 */

#include "KMEnter.h"

#include <sys/capability.h>
#include <filesystem>

#include "common/KMDirFd.h"
#include "common/KMEnv.h"
#include "common/KMInstance.h"
#include "common/KMLogger.h"

using namespace std;
namespace fs = std::filesystem;

class KMEnter::Options : public KMOptions
{
public:
    Options() = default;
    virtual ~Options() = default;

protected:
    /**
     * @brief 命令行参数解析后操作
     */
    void preParseHook() override;

    /**
     * @brief 命令行参数解析后操作
     */
    void postParseHook() override;

private:
    /**
     * @brief 命令行参数添加描述选项
     */
    void addOptions();

    /**
     * @brief 命令行参数检查项
     */
    void checkOptions(const opt::variables_map &varsMap);

public:
    bool m_bHelp;
    string m_instance;
    vector<string> m_cmdArgs;
};

void KMEnter::Options::preParseHook()
{
    addOptions();
}

/**
 * @brief 命令行参数解析后操作
 */
void KMEnter::Options::postParseHook()
{
    checkOptions(m_varsMap);

    if (m_bHelp)
    {
        showUsage();
        exit(EXIT_SUCCESS);
    }
}

/**
 * @brief 命令行参数添加描述选项
 */
void KMEnter::Options::addOptions()
{
    m_optionsDescription = std::make_unique<opt::options_description>(_("\n用法:\n\tkaiming enter instance command - 在正在运行的沙箱中运行命令"));

    opt::options_description positionalOptions(_("Positional Options"));
    positionalOptions.add_options()("instance", opt::value<string>(&m_instance), _("应用实例/应用名称/pid"))("command", opt::value<vector<string>>(&m_cmdArgs), _("容器中运行的命令"));
    opt::options_description helpOptions(_("帮助选项"));
    helpOptions.add_options()("help,h", _("显示帮助选项"));
    m_optionsDescription->add(helpOptions).add(positionalOptions);
    m_optionsPositionalDescription->add("instance", 1);
    m_optionsPositionalDescription->add("command", -1);
}

/**
 * @brief 命令行参数检查项
 */
void KMEnter::Options::checkOptions(const opt::variables_map &varsMap)
{
    checkOption(varsMap, "help", "h", m_bHelp);
}

class KMEnter::Private
{
public:
    int m_pid;
    vector<string> m_vecCommand;
    std::unique_ptr<Options> m_kmOptions;
};

REGISTER_SUBCOMMAND_DYNCREATE(enter, KMEnter)

KMEnter::KMEnter()
    : d(make_unique<Private>())
{
    d->m_kmOptions = make_unique<Options>();
}

KMEnter::~KMEnter()
{
}

/**
 * @brief 子命令运行
 * @param argc 参数个数
 * @param argv 参数数组
 * @return int 0 成功
 */
int KMEnter::despose(int argc, char **argv)
{
    KMDebug("KMEnter::despose invoke begin");

    // 参数解析
    d->m_kmOptions->parseCommandLine(argc, argv);

    // 环境检查
    if (!checkEnv(argc, argv))
    {
        return -1;
    }

    return enter(d->m_pid, d->m_vecCommand) ? 0 : -1;
}

/**
 * @brief 数据解析与环境检查
 * @return true 成功 false 失败
 */
bool KMEnter::checkEnv(int argc, char **argv)
{
    if (argc < 3)
    {
        KMError("INSTANCE and COMMAND must be specified");
        return false;
    }

    // 参数解析
    if (!checkPid(argv[1], d->m_pid))
    {
        KMError(string(argv[1]) + " is neither a pid nor an application or instance ID");
        return false;
    }

    for (int i = 2; i < argc; i++)
        d->m_vecCommand.emplace_back(argv[i]);

    return true;
}

/**
 * @brief enter 子命令运行
 * @param pid 进程pid
 * @param command 子命令参数
 * @return bool true 是 fase 否
 */
bool KMEnter::enter(int &pid, vector<string> &command)
{
    int uid, gid;
    if (!getUidGid(pid, uid, gid))
    {
        KMError("Can't get uid and gid .");
        return false;
    }

    string cwdPath, rootPath;
    if (!getCwdAndRoot(pid, cwdPath, rootPath))
    {
        KMError("Can't get cwd and root path .");
        return false;
    }

    if (!setNameSpace(pid))
    {
        KMError("Set name space fail .");
        return false;
    }

    if (chdir(cwdPath.c_str()))
    {
        KMError("Can't chdir .");
        return false;
    }

    if (chroot(rootPath.c_str()))
    {
        KMError("Can't chroot .");
        return false;
    }

    if (setgid(gid))
    {
        KMError("Can't switch gid .");
        return false;
    }

    if (setuid(uid))
    {
        KMError("Can't switch uid .");
        return false;
    }

    clearAllCaps();

    vector<string> env;
    getEnv(pid, uid, env);

    return spawnSync(command, env);
}

/**
 * @brief 检查并获取子进程的pid
 * @param id instance id
 * @param pid 进程pid
 * @return bool true 是 fase 否
 */
bool KMEnter::checkPid(const string &id, int &pid)
{
    vector<shared_ptr<KMInstance>> instances = KMInstance::getAllInstances();
    for (auto &instance : instances)
    {
        if (instance->getInstanceId().compare(id) == 0 || instance->getApp().compare(id) == 0 || instance->getPid() == atoi(id.c_str()))
        {
            pid = instance->getChildPid();
            return true;
        }
    }

    return false;
}

/**
 * @brief 获取uid gid
 * @param pid 进程id
 * @param uid 用户id
 * @param gid 组id
 * @return bool true 是 fase 否
 */
bool KMEnter::getUidGid(const int &pid, int &uid, int &gid)
{
    string procRobot = string("/proc/") + to_string(pid) + string("/root");
    struct stat statBuf;
    if (stat(procRobot.c_str(), &statBuf) == 0)
    {
        uid = statBuf.st_uid;
        gid = statBuf.st_gid;
        return true;
    }

    if (errno == EACCES)
    {
        KMError("Entering not supported (need unprivileged user namespaces, or sudo -E)");
    }
    else
    {
        KMError("No such pid " + to_string(pid));
    }

    return false;
}

/**
 * @brief 清除进程的所有capability
 */
void KMEnter::clearAllCaps()
{
    // 声明一个用户 capability 头结构体
    struct __user_cap_header_struct hdr = { _LINUX_CAPABILITY_VERSION_3, 0 };
    // 声明一个用户 capability 数据结构体数组，其中包含 2 个元素，每个元素都为 0
    struct __user_cap_data_struct data[2] = { { 0 } };

    // 设置用户 capability，传入头结构体和数据结构体数组
    capset(&hdr, data);
}

/**
 * @brief 获取当前工作目录和根目录
 * @param pid 进程id
 * @param cwdPath 当前工作目录
 * @param rootPath 根目录
 * @return bool true 是 fase 否
 */
bool KMEnter::getCwdAndRoot(int pid, string &cwdPath, string &rootPath)
{
    string cwd = string("/proc/") + to_string(pid) + string("/cwd");
    string root = string("/proc/") + to_string(pid) + string("/root");

    cwdPath = KMDirFd::readlinkat(-1, cwd);
    rootPath = KMDirFd::readlinkat(-1, root);

    return true;
}

/**
 * @brief 设置命名空间
 * @param pid 进程id
 * @return bool true 是 fase 否
 */
bool KMEnter::setNameSpace(int pid)
{
    ino_t userbaseIno = 0;
    const vector<string> ns = { "user_base", "ipc", "net", "pid", "mnt", "user" };
    string root = string("/proc/") + to_string(pid) + string("/root");
    for (auto &name : ns)
    {
        string path, selfpath;
        struct stat pathStat, selfpathStat;

        if (name.compare("user_base") == 0)
        {
            path = root + string("/run/.userns");
            selfpath = string("/proc/self/ns/user");
        }
        else
        {
            path = string("/proc/") + to_string(pid) + string("/ns/") + name;
            selfpath = string("/proc/self/ns/") + name;
        }

        if (stat(path.c_str(), &pathStat) != 0)
        {
            if (errno == ENOENT)
            {
                continue;
            }
            KMError("Invalid " + name + string(" namespace for pid ") + to_string(pid));
            return false;
        }

        if (name.compare("user") == 0 && pathStat.st_ino == userbaseIno)
        {
            continue;
        }

        if (stat(selfpath.c_str(), &selfpathStat) != 0)
        {
            KMError("Invalid " + name + string(" namespace for self"));
            return false;
        }

        if (selfpathStat.st_ino == pathStat.st_ino)
        {
            continue;
        }

        if (name.compare("user_base") == 0)
        {
            userbaseIno = pathStat.st_ino;
        }

        int fd = open(path.c_str(), O_RDONLY);
        if (fd == -1)
        {
            KMError("Can't open " + name + string(" namespace for pid ") + to_string(pid));
            return false;
        }

        if (setns(fd, 0) == -1)
        {
            close(fd);
            KMError("Can't enter " + name + string(" namespace for pid ") + to_string(pid) + " error: " + strerror(errno));
            return false;
        }

        close(fd);
    }

    return true;
}

/**
 * @brief 获取环境变量
 * @param pid 进程id
 * @param uid 用户id
 * @param env 环境变量
 */
void KMEnter::getEnv(int pid, int uid, vector<string> &env)
{
    string envpath = string("/proc/") + to_string(pid) + string("/environ");
    if (fs::exists(fs::path(envpath)))
    {
        do
        {
            string contents;
            KMUtil::readFile(envpath, contents);

            if (contents.empty())
            {
                break;
            }

            if (contents.compare(0, strlen("PULSE_SERVER="), "PULSE_SERVER=") != 0 && contents.compare(0, strlen("PULSE_CLIENTCONFIG="), "PULSE_CLIENTCONFIG=") != 0 &&
                contents.compare(0, strlen("XDG_RUNTIME_DIR="), "XDG_RUNTIME_DIR=") != 0 && contents.compare(0, strlen("DBUS_SYSTEM_BUS_ADDRESS="), "DBUS_SYSTEM_BUS_ADDRESS=") != 0 &&
                contents.compare(0, strlen("DBUS_SESSION_BUS_ADDRESS="), "DBUS_SESSION_BUS_ADDRESS=") != 0)
            {
                env.emplace_back(contents);
            }
        } while (0);
    }

    string xdgRunDirEnv = string("XDG_RUNTIME_DIR=") + string("/run/user/") + to_string(uid);
    env.emplace_back(xdgRunDirEnv);

    string pulsepath = string("/run/user/") + to_string(uid) + string("/pulse/native");
    if (fs::exists(pulsepath))
    {
        env.emplace_back(string("PULSE_SERVER=unix:") + pulsepath);
        env.emplace_back(string("PULSE_CLIENTCONFIG=/run/user/") + to_string(uid) + string("/pulse/config"));
    }

    string sessionBusPath = string("/run/user/") + to_string(uid) + string("/bus");
    if (fs::exists(sessionBusPath))
    {
        env.emplace_back(string("DBUS_SESSION_BUS_ADDRESS=unix:path=") + sessionBusPath);
    }

    string sysBusPath = "/run/dbus/system_bus_socket";
    if (fs::exists(sysBusPath))
    {
        env.emplace_back("DBUS_SYSTEM_BUS_ADDRESS=unix:path=/run/dbus/system_bus_socket");
    }

    if (getenv("TERM"))
    {
        env.emplace_back(string("TERM=") + string(getenv("TERM")));
    }
}

/**
 * @brief 启动子进程
 * @param vArgs 命令行参数
 * @param vEnv 环境变量
 * @return bool true 是 fase 否
 */
bool KMEnter::spawnSync(vector<string> &vArgs, vector<string> &vEnv)
{
    vector<char *> argvarray;
    vector<char *> envparray;

    for (auto &arg : vArgs)
    {
        argvarray.push_back(const_cast<char *>(arg.c_str()));
    }
    argvarray.push_back(nullptr);

    for (auto &arg : vEnv)
    {
        envparray.push_back(const_cast<char *>(arg.c_str()));
    }
    envparray.push_back(nullptr);

    signal(SIGINT, SIG_IGN);

    if (execvpe(argvarray[0], (char **)(argvarray.data()), (char **)envparray.data()) != 0)
    {
        return false;
    }

    return true;
}