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

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

# if you using Linux based OS, please run this command: mount /dev/sdX /mnt/disk --mkdir
# DISK = "/mnt/disk"


def find_disk():
    if platform.system() == 'Linux':
        if "{disk}" in DISK:
            raise OSError(
                "Cannot auto find the Ventoy disk in Linux! Please change \"DISK\"")
    elif platform.system() == 'Windows':
        ventoy_drives = []
        for i in "CDEFGHIJKABLMNOPQRSTUVWYZ":
            this_disk = i+":"
            if os.path.exists(this_disk + "/ventoy/ventoy.json"):
                ventoy_drives.append(this_disk)
        if len(ventoy_drives) != 1:
            for (index, value) in enumerate(ventoy_drives):
                print(f"{i}:{ventoy_drives[i]}", end=" ")
            return ventoy_drives[int(input("Please choose your Ventoy device:"))]
    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("Making 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.")