151 lines
4.4 KiB
Python
151 lines
4.4 KiB
Python
#####################################
|
|
## 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.")
|