peinjector/installer/install.py

#####################################
##     PEinjector/installer        ##
#####################################

# This moudle is a part of PEinjector.
# Please comply with the LICENSE.

import os
import shutil
import json
import platform
import string

############ Config ############
KEYWORD = "winpe"
PYTHON = "python"
DISK = "{disk}"
################################


def find_disk():
    ventoy_drives = []
    if platform.system() == 'Linux':
        mount_point = "/mnt/prinjector_mount_point"
        if not os.path.exists(mount_point):
            os.mkdir(mount_point)
        elif len(os.listdir(mount_point)):
            raise OSError("Mount point {} not empty!".format(mount_point))

        i: int = 97
        while os.path.exists(f"/dev/sd{chr(i)}"):
            j: int = 1
            this_disk = f"/dev/sd{chr(i)}{str(j)}"
            while os.path.exists(this_disk):
                os.system(f"mount /dev/sd{chr(i)}{str(j)} {mount_point}")
                if os.path.exists(mount_point+"/ventoy/ventoy.json"):
                    return mount_point
                j += 1
                os.system(f"umount {mount_point}")
            i += 1

        i: int = 0
        while os.path.exists(f"/dev/nvme{str(i)}"):
            j: int = 1
            while os.path.exists(f"/dev/nvme{str(i)}n{str(j)}"):
                k: int = 1
                while os.path.exists(f"/dev/nvme{str(i)}n{str(j)}p{str(k)}"):
                    os.system(
                        f"mount /dev/nvme{str(i)}n{str(j)}p{str(k)} {mount_point}")
                    if os.path.exists(mount_point+"/ventoy/ventoy.json"):
                        return mount_point
                    k += 1
                    os.system(f"umount {mount_point}")
                j += 1
            i += 1

    elif platform.system() == 'Windows':
        for i in "CDEFGHIJKABLMNOPQRSTUVWYZ":
            this_disk = i+":"
            if os.path.exists(this_disk + "/ventoy/ventoy.json"):
                return this_disk

    else:
        raise OSError(f"Unsupported operating system:{platform.system()}")
    raise OSError(
        "Cannot find the Ventoy disk!(try to run ventoyPlugson once)")


DISK = DISK.replace("{disk}", find_disk())

print("Use disk: "+DISK)
with open(DISK+"/ventoy/ventoy.json", "r") as file:
    ventoy_cfg = json.load(file)
searchroot = ""
search_level = 3
if "control" in ventoy_cfg:
    for i in ventoy_cfg["control"]:
        if "VTOY_DEFAULT_SEARCH_ROOT" in i:
            searchroot = i["VTOY_DEFAULT_SEARCH_ROOT"]
        if "VTOY_MAX_SEARCH_LEVEL" in i:
            if i["VTOY_MAX_SEARCH_LEVEL"] != "max":
                search_level = int(i["VTOY_MAX_SEARCH_LEVEL"])
filelist = []


def dfs(deep, dir, search_level):
    if deep > search_level:
        return
    lists = os.listdir(DISK+dir)
    for i in lists:
        if os.path.isfile(DISK+dir+"/"+i):
            if ".iso" in i.lower() and KEYWORD in i.lower():
                filelist.append(dir+"/"+i)
        else:
            dfs(deep+1, dir+"/"+i, search_level)


dfs(0, searchroot, search_level)

if "injection" in ventoy_cfg:
    for i in ventoy_cfg["injection"]:
        if "image" in i:
            if i["image"] in filelist:
                filelist.remove(i["image"])
else:
    ventoy_cfg["injection"] = []

print("Found iso:", filelist)
for i in filelist:
    ventoy_cfg["injection"].append({
        "image": i,
        "archive": "/PEinjector/PEinjector.7z"
    })


print("Writing config")
with open(DISK+"/ventoy/ventoy.json", "w") as file:
    file.write(json.dumps(ventoy_cfg))


print("Copying files")
if not os.path.exists(DISK+"/PEinjector"):
    os.mkdir(DISK+"/PEinjector")
if os.path.exists(DISK+"/PEinjector/PEinjector.7z"):
    os.remove(DISK+"/PEinjector/PEinjector.7z")
shutil.copyfile("dist/PEinjector.7z", DISK+"/PEinjector/PEinjector.7z")
if not os.path.exists(DISK+"/PEinjector/disable.txt"):
    with open(DISK+"/PEinjector/disable.txt", "w") as file:
        pass


print("Writing version info")
shutil.copyfile("version", DISK+"/PEinjector/VERSION")


print("Creating dir")
if not os.path.exists(DISK+"/PEinjector/install"):
    os.mkdir(DISK+"/PEinjector/install")


print("Installing basic package")
if not os.path.exists(DISK+"/PEinjector/package"):
    os.mkdir(DISK+"/PEinjector/package")
for i in os.listdir("package"):
    if os.path.exists(DISK+"/PEinjector/package/"+i):
        shutil.rmtree(DISK+"/PEinjector/package/"+i)
    shutil.copytree("package/"+i, DISK+"/PEinjector/package/"+i)


print("All done.")