injectorHub/lib/get_disk.py

import platform
import os
import subprocess
import sys

from ui.lang import l
from logger.logger import *

if platform.system() == 'Linux':
    def find_hub_disk():
        info("[getdisk]get hub disk")
        output = subprocess.check_output(
            ['mount', '-v']).decode(sys.getdefaultencoding())
        for line in output.split('\n'):
            if not line:
                continue
            mount_point_l = line.split(" ")
            disk_id = mount_point_l[0]
            mount_to = mount_point_l[2]
            if len(disk_id) <= len("/dev/sd") or (disk_id[:len("/dev/sd")] != "/dev/sd" and (disk_id[:len("/dev/sd")] != "/dev/lo") and (disk_id[:len("/dev/sd")] != "/dev/nv")):
                continue
            if (mount_to == "/" or mount_to == "/boot"):
                continue
            if (os.path.exists(mount_to + "/PEinjector")):
                return (0, mount_to, disk_id)
        return (2, "")

    def find_vtoy_disk():
        vtoylist = []
        info("[getdisk]get vtoy disk")
        output = subprocess.check_output(
            ['mount', '-v']).decode(sys.getdefaultencoding())
        for line in output.split('\n'):
            if not line:
                continue
            mount_point_l = line.split(" ")
            disk_id = mount_point_l[0]
            mount_to = mount_point_l[2]
            if len(disk_id) <= len("/dev/sd") or (disk_id[:len("/dev/sd")] != "/dev/sd" and (disk_id[:len("/dev/sd")] != "/dev/lo") and (disk_id[:len("/dev/sd")] != "/dev/nv")):
                continue
            if (mount_to == "/" or mount_to == "/boot"):
                continue
            if os.path.exists(mount_to + "/ventoy"):
                vtoylist.append((disk_id, mount_to))
                continue
            output2 = subprocess.check_output(
                ['blkid', '-o', 'value', '-s', 'LABEL', disk_id]).decode(sys.getdefaultencoding())
            output2 = output2.replace('\n', "")
            if (output2 == ""):
                continue
            if (output2 == "Ventoy" or output2 == "ventoy"):
                vtoylist.append((disk_id, mount_to))
        return vtoylist

    def find_all_disk():
        vtoylist = []
        info("[getdisk]get all disk")
        output = subprocess.check_output(
            ['mount', '-v']).decode(sys.getdefaultencoding())
        for line in output.split('\n'):
            if not line:
                continue
            mount_point_l = line.split(" ")
            disk_id = mount_point_l[0]
            mount_to = mount_point_l[2]
            if len(disk_id) <= len("/dev/sd") or (disk_id[:len("/dev/sd")] != "/dev/sd" and (disk_id[:len("/dev/sd")] != "/dev/lo") and (disk_id[:len("/dev/sd")] != "/dev/nv")):
                continue
            if (mount_to == "/" or mount_to == "/boot"):
                continue
            vtoylist.append((disk_id, mount_to))
        return vtoylist
elif platform.system() == 'Windows':
    def find_hub_disk():
        info("[getdisk]get hub disk")
        for i in "CDEFGHIJKABLMNOPQRSTUVWYZ":
            this_disk = i+":"
            if os.path.exists(this_disk + "/PEinjector"):
                return (0, this_disk, this_disk)
        return (2, "")

    def find_vtoy_disk():
        vtoylist = []
        info("[getdisk]get vtoy disk")
        for i in "CDEFGHIJKABLMNOPQRSTUVWYZ":
            this_disk = i+":"
            if not os.path.exists(this_disk):
                continue
            if os.path.exists(this_disk + "/ventoy"):
                vtoylist.append((this_disk, this_disk))
            with os.popen("wmic logicaldisk" +
                          f" where name=\"{this_disk}\" get volumename", 'r', 1) as file:
                output2 = file.read()
            output2 = (output2.replace("\n", "")
                       .replace("VolumeName", "")
                       .lstrip()
                       .rstrip())
            if (output2 == "Ventoy"):
                vtoylist.append((this_disk, this_disk))
        return vtoylist

    def find_all_disk():
        vtoylist = []
        info("[getdisk]get all disk")
        for i in "CDEFGHIJKABLMNOPQRSTUVWYZ":
            this_disk = i+":"
            if (os.path.exists(this_disk)):
                vtoylist.append((this_disk, this_disk))
        return vtoylist
else:
    def find_hub_disk():
        info("[getdisk]get hub disk")
        return (1, l("UnsupportedSystem"))