#!/usr/bin/env python
# coding:utf-8
# Copyright (c) 2024 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""
文件名：prototype.py
功能：默认裁剪脚本
"""
import os
import sys
import json
import stat
import logging
 
cwd = os.path.join(os.getcwd(), os.path.dirname(__file__))
sys.path.append(f"{cwd}/../")
sys.path.append(f"{cwd}/../../")
sys.path.append(f"{cwd}/../../../")
from customization.default import BaseCustomization


log = logging.getLogger(__name__)

cut_list = [
    # /etc
    'etc/init.d/README',
    'etc/kernel/', # 空文件夹
    'etc/resolv-conf.systemd',
    'etc/X11/xinit/xinitrc.d/50-systemd-user.sh',
    'etc/mdev/find-touchscreen.sh',
    'etc/shadow-maint/', # shadow开源脚本文件夹，userdel时调用

    # /lib
    'lib/modprobe.d/README',
    'lib/systemd/system/auditd.service',
    'lib/systemd/system-preset/98-auditd.preset',

    # /bin
    'bin/chattr.e2fsprogs',

    # /sbin
    'sbin/mke2fs.e2fsprogs',
    'sbin/vipw',

    # /usr/bin
    'usr/bin/attr',
    'usr/bin/chage',
    'usr/bin/chfn',
    'usr/bin/chsh',
    'usr/bin/expiry',
    'usr/bin/faillog',
    'usr/bin/getfattr',
    'usr/bin/kernel-install',
    'usr/bin/lastlog',
    'usr/bin/lsattr',
    'usr/bin/newgidmap',
    'usr/bin/newuidmap',
    'usr/bin/passwd',
    'usr/bin/pwdx',
    'usr/bin/pwdx.procps',
    'usr/bin/setfattr',
    'usr/bin/skill',
    'usr/bin/snice',
    'usr/bin/tput',
    'usr/bin/tset',
    'usr/bin/w',
    'usr/bin/w.procps',
    'usr/bin/xmlcatalog',
    'usr/bin/xmllint',

    # /usr/lib64
    'usr/lib64/busybox', # 这些脚本全都只调用/bin/busybox.nosuid

    # /usr/sbin
    'usr/sbin/chpasswd',
    'usr/sbin/groupadd',
    'usr/sbin/groupdel',
    'usr/sbin/groupmems',
    'usr/sbin/groupmod',
    'usr/sbin/grpck',
    'usr/sbin/grpconv',
    'usr/sbin/grpunconv',
    'usr/sbin/jffs2dump',
    'usr/sbin/jffs2reader',
    'usr/sbin/newusers',
    'usr/sbin/pwck',
    'usr/sbin/pwconv',
    'usr/sbin/pwunconv',
    'usr/sbin/sumtool',
    'usr/sbin/useradd',
    'usr/sbin/userdel',
    'usr/sbin/usermod'

    # /usr/share
    'usr/share/terminfo/a/ansi',
    'usr/share/terminfo/r/rxvt',
    'usr/share/terminfo/s/screen-256color',
    'usr/share/terminfo/s/sun',
    'usr/share/terminfo/v/vt220',
    'usr/share/terminfo/v/vt52',
    'usr/share/terminfo/x/xterm-color',
    'usr/share/terminfo/x/xterm-256color',
    'usr/share/terminfo/x/xterm-xfree86',

    # systemd
    'lib/systemd/system/sysinit.target.wants/dev-hugepages.mount',
    'lib/systemd/system/dev-hugepages.mount',
    'lib/systemd/system/sysinit.target.wants/dev-mqueue.mount',
    'lib/systemd/system/dev-mqueue.mount',
    'lib/systemd/system/serial-getty@.service'
]


subsys_startup_components = [
    # bmc_core components
    'firmware_mgmt',
    'bmc_upgrade',
    'bmc_time',
    'fructrl',
    'bmc_network',
    'bmc_health',
    'ipmi_core',
    'bmc_soc',
    # alarm components
    'sensor',
    'host_agent',
    'frudata',
    'event',
    # framework components
    'libsoc_adapter',
    'hwproxy',
    'hwdiscovery',
    'soctrl',
    'maca',
    "key_mgmt",
    "persistence",
    # hardware components
    'manufacture',
    'bios',
    'general_hardware',
    'thermal_mgmt',
    'power_mgmt',
    'pcie_device',
    'chassis',
    'network_adapter',
    'storage',
    'compute',
    'mctpd',
    # interface components
    'redfish',
    'web_backend',
    # om_priv components
    'nsm',
    'oms',
    # om components
    'ddns',
    'remote_console',
    'event_policy',
    'usb_entry',
    'product_mgmt',
    'rmcpd',
    # security components
    'iam',
    'certificate'
]


class Customization(BaseCustomization):
    def collect_table_info(self, dir_name, table_info):
        service_file = os.path.join(dir_name, "mds/service.json")
        if not os.path.exists(service_file):
            return
        
        app_name = ""
        with open(service_file, "r") as service_fp:
            service_data = json.load(service_fp)
            if "name" not in service_data:
                return
            app_name = service_data["name"]

        model_file = os.path.join(dir_name, "mds/model.json")
        if not os.path.exists(model_file):
            return

        with open(model_file, "r") as model_fp:
            model_data = json.load(model_fp)
            for _, class_data in model_data.items():
                if "tableName" in class_data and class_data.get("tableLocation", "") != "Local":
                    table_name = class_data["tableName"]
                    table_info[table_name] = table_info.get(table_name, [])
                    table_info[table_name].append(app_name)
                    table_info[table_name].sort()

    def package_table_owner(self, rootfs_path):
        conan_install = os.path.join(self.config.build_path, "conan_install")
        if os.path.exists(conan_install) is False:
            raise FileExistsError(f"{conan_install} 文件夹不存在, 无法打包文件")

        table_info = {}
        for root, dirs, _ in os.walk(conan_install):
            dirs.sort()
            for dire in dirs:
                dir_name = os.path.join(root, dire)
                if "include" in dir_name and os.path.isdir(dir_name):
                    self.collect_table_info(dir_name, table_info)

        temp_table_owner_dir = os.path.join(self.config.build_path, "temp_table_owner")
        os.makedirs(temp_table_owner_dir, exist_ok=True)
        temp_table_owner_file = os.path.join(temp_table_owner_dir, "table_owner.json")
        flags = os.O_WRONLY | os.O_CREAT | os.O_TRUNC
        mode = stat.S_IWUSR | stat.S_IRUSR
        with os.fdopen(os.open(temp_table_owner_file, flags, mode), 'w+') as table_owner_fp:
            json.dump(table_info, table_owner_fp)

        table_owner_file = f"{rootfs_path}/opt/bmc/apps/persistence/mds/table_owner.json"
        self.work.run_command(f"cp {temp_table_owner_file} {table_owner_file}", sudo=True)
        self.work.run_command(f"chown root:root {table_owner_file}", sudo=True)
        self.work.run_command(f"chmod 640 {table_owner_file}", sudo=True)

    def get_systemd_name(self, process_name):
        if process_name == "ssdp":
            return "bmc_" + process_name
        
        return process_name

    def cutout_systemd_conf(self, rootfs_path, process_name):
        systemd_name = self.get_systemd_name(process_name)
        systemd_config_path = os.path.join(f"{rootfs_path}/etc/systemd/system", systemd_name + ".service")
        if os.path.exists(systemd_config_path) is False:
            systemd_config_path = os.path.join(f"{rootfs_path}/lib/systemd/system", systemd_name + ".service")
            if os.path.exists(systemd_config_path) is False:
                return

        self.work.run_command(f"rm -rf {systemd_config_path}", sudo=True)

    def cutout_systemd_confs(self, rootfs_path):
        launch_control_json = f"{rootfs_path}/opt/bmc/conf/launch_control.json"
        if os.path.exists(launch_control_json) is False:
            return
        
        self.work.run_command(f"sudo chown root:root {launch_control_json}", sudo=True)
        self.work.run_command(f"sudo chmod 444 {launch_control_json}", sudo=True)
        with open(launch_control_json) as f:
            config_data = json.load(f)
            for process_name, process_config in config_data.items():
                if "systemd_conf" not in process_config:
                    self.cutout_systemd_conf(rootfs_path, process_name)

    def rootfs_common(self, rootfs_path):
        log.info("---->>> Start run qemu customization configure, path:{} <<<----".format(rootfs_path))
        cur_dir = os.getcwd()
        os.chdir(rootfs_path)
        self.package_table_owner(rootfs_path)
        self.cutout_systemd_confs(rootfs_path)
        self.work.run_command("rm -rf usr/lib64/mock", sudo=True)
        os.chdir(cur_dir)

    def create_memory_monitor_confs(self, rootfs_path):
        launch_control_json = f"{rootfs_path}/opt/bmc/conf/launch_control.json"
        if os.path.exists(launch_control_json) is False:
            return

        mem_conf_file = f"{rootfs_path}/opt/bmc/conf/process_config.ini"
        if os.path.exists(mem_conf_file) is False:
            self.work.run_command(f"touch {mem_conf_file}", sudo=True)

        self.work.run_command(f"sudo chown root:root {launch_control_json}", sudo=True)
        self.work.run_command(f"sudo chmod 444 {launch_control_json}", sudo=True)
        self.work.run_command(f"sudo chown root:root {mem_conf_file}", sudo=True)
        self.work.run_command(f"sudo chmod 666 {mem_conf_file}", sudo=True)
        ini_f = open(mem_conf_file, 'w')
        with open(launch_control_json) as f:
            config_data = json.load(f)
            for process, process_config in config_data.items():
                if process == "processes":
                    self.update_memory_config(ini_f, process_config)

        self.work.run_command(f"sudo chmod 444 {mem_conf_file}", sudo=True)

    def update_memory_config(self, ini_f, process_config):
        for process_name, config in process_config.items():
            if "mem_max_kb" in config:
                mem_max = config["mem_max_kb"]
                ini_f.write(f"{process_name} {mem_max}\n")

    def cutout_subsys_startup_comps(self, rootfs_path):
        """
        对于需要合并到子系统拉起的组件，需要对组件自带的systemd配置文件进行裁剪，否则可能出现组件被多次拉起的问题，
        在subsys_startup_components中添加组件名即可完成对组件自带的systemd配置文件的裁剪
        """
        for component in subsys_startup_components:
            path = f"{rootfs_path}/etc/systemd/system/{component}.service"
            self.work.run_command(f"rm -rf {path}", sudo=True)
            multi_user_path = f"{rootfs_path}/etc/systemd/system/multi-user.target.wants/{component}.service"
            self.work.run_command(f"rm -rf {multi_user_path}", sudo=True)
        self.work.run_command(f"rm -rf {rootfs_path}/opt/bmc/apps/mdb_mgmt", sudo=True)

    def rootfs_cust(self, rootfs_path):
        log.info("---->>> Start run prototype customization configure, path:{} <<<----".format(rootfs_path))
        os.chdir(rootfs_path)
        log.info("change directory:{}".format(rootfs_path))
        self.work.run_command("sed -i '5i After=rc-local.service' lib/systemd/system/dbus.service", sudo=True)
        self.work.run_command("sed -i 's/messagebus/root/g' usr/share/dbus-1/system.conf", sudo=True)
        self.rootfs_service_cust()

        for path in cut_list:
            self.work.run_command(f"rm -rf {path}", sudo=True)

        # 项目的 conan 文件里面不能创建目录的软连接，否则会打包失败，只能放在这里了
        self.work.pipe_command(["sudo find etc/rc.d/rc.start -mindepth 1", "sudo xargs -P 0 chmod 550"])
        self.work.run_command("rm -rf home", sudo=True)
        self.work.run_command("ln -s /data/trust/home home", sudo=True)

        # 删除logrotate配置文件
        self.work.run_command("rm -rf etc/logrotate.d/btmp", sudo=True)
        self.work.run_command("rm -rf etc/logrotate.d/wtmp", sudo=True)

        # 依赖包权限修改
        self.work.run_command("sudo chmod 444 etc/package_info", sudo=True)
        self.work.run_command("sudo chmod 444 opt/bmc/trust/sensitive_data.json", sudo=True)
        self.work.run_command("sudo chmod 711 bin/busybox.nosuid", sudo=True)

        # CLI命令别名设置
        self.work.pipe_command([
            "echo \"\nif [ -f /opt/bmc/script/cpu_monitor.sh ]; then\"",
            "sudo tee -a etc/rc.d/rc.sysinit"
        ])
        self.work.pipe_command(["echo \"    sh /opt/bmc/script/cpu_monitor.sh &\"", "sudo tee -a etc/rc.d/rc.sysinit"])
        self.work.pipe_command(["echo \"fi\n\"", "sudo tee -a etc/rc.d/rc.sysinit"])
        self.work.pipe_command([
            "echo \"alias ipmcget='/opt/bmc/skynet/lua /opt/bmc/apps/cli/service/ipmcget.lua'\"",
            f"sudo tee -a \"{rootfs_path}/etc/profile\""
        ])
        self.work.pipe_command([
            "echo \"alias ipmcset='/opt/bmc/skynet/lua /opt/bmc/apps/cli/service/ipmcset.lua'\"",
            f"sudo tee -a \"{rootfs_path}/etc/profile\""
        ])

        self.work.run_command("mkdir -p data/opt/bmc/web/custom", sudo=True)
        self.work.run_command("sudo chmod 750 data/opt/bmc/web/custom", sudo=True)
        self.create_memory_monitor_confs(rootfs_path)

        self.cutout_subsys_startup_comps(rootfs_path)

        ## 删除snmp中不必要的mib文件和readme
        snmp_mib_path = os.path.join(rootfs_path,
            "opt/bmc/apps/snmp/interface_config/mib/HUAWEI-SERVER-iBMC-MIB.mib")
        self.work.run_command(f"rm -rf {snmp_mib_path}", sudo=True)
        snmp_openubmc_mib_path = os.path.join(rootfs_path,
            "opt/bmc/apps/snmp/interface_config/mib/OPENUBMC-SERVER-MIB.mib")
        self.work.run_command(f"rm -rf {snmp_openubmc_mib_path}", sudo=True)
        snmp_mib_read_me_path = os.path.join(rootfs_path, "opt/bmc/apps/snmp/interface_config/mib/README.md")
        self.work.run_command(f"rm -rf {snmp_mib_read_me_path}", sudo=True)

    def rootfs_service_cust(self):
        self.work.run_command(
            "sed -i '$a ExecStartPost=/opt/bmc/rundbus.sh' lib/systemd/system/dbus.service",
            sudo=True
        )
        self.work.run_command(
            "sed -i -E 's/(^After=.*$)/\\0 rc-local.service/' "
            "lib/systemd/system/systemd-user-sessions.service",
            sudo=True
        )
        # systemd-user-sessions服务启动失败时需要重试
        self.work.run_command(
            "sed -i '/RemainAfterExit/i Restart=on-failure' lib/systemd/system/systemd-user-sessions.service", 
            sudo=True
        )
        self.work.run_command("mkdir -p etc/systemd/system/multi-user.target.wants", sudo=True)
        self.work.run_command(
            "ln -s /lib/systemd/system/dbus.service etc/systemd/system/multi-user.target.wants/dbus.service",
            sudo=True
        )

        # 裁剪
        self.work.run_command("mkdir -p run", sudo=True)
        self.work.run_command("mkdir -p var", sudo=True)
        self.work.run_command("rm sbin/init -f", sudo=True)
        self.work.run_command("ln -s ../lib/systemd/systemd sbin/init", sudo=True)
        self.work.run_command(
            "ln -s /lib/systemd/system/multi-user.target etc/systemd/system/default.target",
            sudo=True
        )
        self.work.run_command(
            "ln -s /lib/systemd/system/getty@.service etc/systemd/system/getty.target.wants/getty@ttyS0.service",
            sudo=True
        )
        self.work.run_command("ln -s /etc/rc.d/rc.sysinit etc/rc.local", sudo=True)
        self.work.run_command(
            "ln -s /lib/systemd/system/rc-local.service etc/systemd/system/rc-local.service",
            sudo=True
        )
        self.work.run_command(
            "mv etc/systemd/system/getty.target.wants/serial-getty@ttyAMA0.service "
            "etc/systemd/system/getty.target.wants/serial-getty@ttyS0.service", 
            sudo=True
        )
        self.work.run_command(
            "ln -s /lib/systemd/system/sshd@.service etc/systemd/system/multi-user.target.wants/sshd@.service",
            sudo=True
        )

        self.work.run_command("sed -i \"/insmod bt_drv.ko/d\" lib/modules/insmod_drv.sh", sudo=True)
        self.work.pipe_command([
            "echo \"/lib/modules/insmod_drv.sh /lib/modules/ko /usr/bin /lib/modules\"",
            f"sudo tee -a \"etc/rc.d/rc.local\""
        ])
        self.work.pipe_command(["echo \"mount -o rw,remount /\"", "sudo tee -a \"etc/rc.d/rc.local\""])

    def datafs_cust(self, rootfs_path):
        pass

    def rootfs_debug_cust(self, rootfs_path):
        os.chdir(rootfs_path)
        self.work.run_command("ln -sf /bin/busybox.nosuid usr/bin/df", sudo=True)
        self.work.pipe_command(["echo \"mkdir -p /data/var/coredump/\"", "sudo tee -a \"etc/rc.d/rc.sysinit\""])
        strip_path = os.path.join(self.config.cross_compile_install_path, "bin", self.config.strip)
        self.work.run_command(f'{strip_path} {rootfs_path}/opt/bmc/libmc/luaclib/dbus.so', sudo=True)

    def rootfs_release_cust(self, rootfs_path):
        os.chdir(rootfs_path)
        self.work.run_command("sed -i '/export \$(cat \/dev\/shm\/dbus\/\.dbus)/d' etc/profile", sudo=True)
        self.work.run_command("ln -sf /bin/busybox.nosuid usr/bin/df", sudo=True)
        self.work.run_command("rm -rf usr/bin/busctl", sudo=True)
        self.work.run_command("rm -rf usr/bin/top+", sudo=True)
        self.work.run_command("rm -rf usr/bin/lsof", sudo=True)
        self.work.run_command("rm -rf usr/bin/traceroute", sudo=True)
        self.work.run_command("rm -rf usr/bin/traceroute6", sudo=True)

        for file in os.listdir("etc/systemd/system/"):
            if not file.endswith(".service"):
                continue
            service_file = os.path.join("etc/systemd/system/", file)
            if not os.path.isfile(service_file):
                continue
            self.work.run_command("sed -i '/LimitSTACK/d' " + service_file, sudo=True)
            self.work.run_command("sed -i '/LimitCORE/d' " + service_file, sudo=True)
