#!/usr/bin/env python3
import os
import sys
import pycdlib
import subprocess


class LiveISO:
    m_iso = pycdlib.PyCdlib()
    m_vmlinuzInitrdSearchPath = "/LIVE/"
    def __init__(self, isoPath):
        self.m_iso.open(isoPath)

    def showList(self, iso_path: str):
        file_list = []
        for i in self.m_iso.list_children(iso_path=iso_path):
            name = str(i.file_identifier(), encoding = "utf-8")
            if (name == ".." or
                name == "."):
                continue
            file_list.append(name)
        return file_list
    
    def vmlinuzPath(self):
        liveList = self.showList(self.m_vmlinuzInitrdSearchPath)
        for i in liveList:
            if ("VMLINU" in i):
                return self.m_vmlinuzInitrdSearchPath + i
        return None
    
    def initrdPath(self):
        liveList = self.showList(self.m_vmlinuzInitrdSearchPath)
        for i in liveList:
            if ("INITRD" in i):
                return self.m_vmlinuzInitrdSearchPath + i
        return None
    
    def filesystemPath(self):
        liveList = self.showList(self.m_vmlinuzInitrdSearchPath)
        for i in liveList:
            if ("FILESYSTEM" in i):
                return self.m_vmlinuzInitrdSearchPath + i
        return None
            
    def extractVmlinuz(self, extractPath):
        vmlinuzPath = self.vmlinuzPath()
        self.extractFile(vmlinuzPath, extractPath)
        

    def extractInitrd(self, extractPath):
        initrdPath = self.initrdPath()
        self.extractFile(initrdPath, extractPath)

    def extractFilesystem(self, extractPath):
        filesystemPath = self.filesystemPath()
        self.extractFile(filesystemPath, extractPath)

    def extractFile(self, iso_path, local_path):
        file = open(local_path, "wb+")
        # 分块读取写入以减少内存使用
        with self.m_iso.open_file_from_iso(iso_path=iso_path) as infp:
            while True:
                data = infp.read(10 * 1024 * 1024)
                if (data == b''):
                    break
                file.write(data)
        file.close()

def RootChecked():
    if (subprocess.getoutput("whoami") != "root"):
        print("请使用 root 运行或者使用 sudo")
        exit(1)

def ShowHelp():
    print(f"{sys.argv[0]} set GXDE OS ISO 路径")
    print("    从 ISO 提取 live 镜像")
    print(f"{sys.argv[0]} del")
    print("    移除安装到本地的 live 镜像")
    print(f"{sys.argv[0]} check")
    print("    检测是否已经安装到本地")

def UpdateGrub():
    os.system("update-grub")

def CheckISOErrorPrint():
    print("No!")
    print("暂不支持该 ISO 镜像")
    exit(1)

if (len(sys.argv) <= 1 or 
    (sys.argv[1] == "set" and len(sys.argv) <= 2)):
    print("参数不足")
    ShowHelp()
    exit(1)

if ("--help" in sys.argv):
    ShowHelp()
    exit()

bootInstallPath = "/boot/gxde/"
liveDiskInstallPath = "/gxde-recovery-live/"

if (sys.argv[1] == "set"):
    # 权限检测
    RootChecked()

    livePath = sys.argv[2]
    # 新建所需文件夹（如果不存在）
    if (not os.path.exists(bootInstallPath)):
        os.makedirs(bootInstallPath)
    if (not os.path.exists(liveDiskInstallPath)):
        os.makedirs(liveDiskInstallPath)

    iso = LiveISO(sys.argv[2])
    # 检测文件是否齐全
    print("检测是否有 vmlinuz？  ", end="")
    if(iso.vmlinuzPath() == None):
        CheckISOErrorPrint()
    else:
        print("Yes")
    print("检测是否有 initrd？  ", end="")
    if(iso.initrdPath() == None):
        CheckISOErrorPrint()
    else:
        print("Yes")
    print("检测是否有 filesystm.squashfs？  ", end="")
    if(iso.filesystemPath() == None):
        CheckISOErrorPrint()
    else:
        print("Yes")


    print("提取 vmlinuz")
    iso.extractVmlinuz(f"{bootInstallPath}/vmlinuz-gxde")
    print("提取 initrd.img")
    iso.extractInitrd(f"{bootInstallPath}/initrd.img-gxde")
    print("提取 filesystem.squashfs")
    iso.extractFilesystem(f"{liveDiskInstallPath}/filesystem.squashfs")
    print("提取完成！刷新 grub")
    UpdateGrub()
    print("完成！重启后即可在 grub 寻找入口")

elif (sys.argv[1] == "del"):
    # 权限检测
    RootChecked()

    if (not os.path.exists(f"{liveDiskInstallPath}/filesystem.squashfs")):
        print("未安装 live 镜像，不进行操作")
        exit(1)
    # 移除镜像
    if (os.path.exists(f"{bootInstallPath}/vmlinuz-gxde")):
        print("移除 vmlinuz")
        os.remove(f"{bootInstallPath}/vmlinuz-gxde")
    if (os.path.exists(f"{bootInstallPath}/initrd.img-gxde")):
        print("移除 initrd.img")
        os.remove(f"{bootInstallPath}/initrd.img-gxde")
    if (os.path.exists(f"{liveDiskInstallPath}/filesystem.squashfs")):
        print("移除 filesystem.squashfs")
        os.remove(f"{liveDiskInstallPath}/filesystem.squashfs")
    print("移除完成！刷新 grub")
    UpdateGrub()
    print("移除完成！")


elif (sys.argv[1] == "check"):
    if (os.path.exists(f"{liveDiskInstallPath}/filesystem.squashfs")):
        print("true")
        exit()
    else:
        print("false")
        exit(1)

else:
    print(f"不支持参数 {sys.argv[1]}")
    exit(1)