import xml.etree.ElementTree as ET
import psutil
import math
import subprocess

def getSize(bytes: int):
    if bytes == 0:
        return {"str": "0B", "int": 0, "name": "B"}
    size_name = ("B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB")
    i = int(math.floor(math.log(bytes, 1024)))
    p = math.pow(1024, i)
    s = round(bytes / p, 2)
    return {"str": f"{s} {size_name[i]}", "int": s, "name": size_name[i]}

class Partition:
    def __init__(self):
        self.mount_point: str = None
        self.capacity: str = None
        self.used: str = None
        self.free: str = None
        self.bitlocker: str = None

class Drive:
    def __init__(self):
        self.model: str = None
        self.sn: str = None
        self.capacity: str = None
        self.paritions: list[Partition] = []

class Drives:
    def __init__(self):
        self.nvme: list[Drive] = []
        self.sata: list[Drive] = []

class GPU:
    def __init__(self):
        self.name: str = None
        self.vram: str = None
        self.clock: str = None
        self.m_clock: str = None

class RAMStick:
    def __init__(self):
        self.mod_num: int = 0
        self.size: str = None
        self.mem_type: str = None
        self.speed: str = None
        self.brand: str = None
        self.pn: str = None

class RAM:
    def __init__(self):
        self.total: str = None
        self.sticks: list[RAMStick] = []

class Motherboard:
    def __init__(self):
        self.model: str = None
        self.chipset: str = None
        self.usb_version: str = None
        self.tpm: str = None

class CPU:
    def __init__(self):
        self.name: str = None
        self.cores: int = 0
        self.lcores: int = 0
        self.hz: str = None

class System:
    def __init__(self):
        self.mobo: Motherboard = Motherboard()
        self.gpu: GPU = GPU()
        self.ram: RAM = RAM()
        self.cpu: CPU = CPU()
        self.drives: Drives = Drives()
        self.name: str = None
        self.os: str = None
        self.brand: str = None
        self.uefi: str = None
        self.sb: str = None

system = System()

def wait_close():
    input("Press any key to close.")
    exit(1)

#tree = ET.parse("DESKTOP-1OS8C68.XML")
tree = ET.parse("Desktop-T493T.XML")

root = tree.getroot()

sub_nodes = root[0].find("SubNodes")

if sub_nodes is None:
    print("Can't find SubNodes element!")
    wait_close()

computer_node = root.find("COMPUTER")
for c in computer_node:
    if c.tag == "Property":
        d = c.find("Description").text
        match c.find("Entry").text:
            case "Computer Name":
                system.name = d
        # system.brand = d
            case "Operating System":
                system.os = d
            case "UEFI Boot":
                system.uefi = d
            case "Secure Boot":
                system.sb = d

mobo_node = sub_nodes.find("MOBO")
for m in mobo_node:
    if m.tag == "Property":
        d = m.find("Description").text
        match m.find("Entry").text:
            case "Motherboard Model":
                system.mobo.model = d
            case "Motherboard Chipset":
                system.mobo.chipset = d
            case "USB Version Supported":
                system.mobo.usb_version = d
            case "Trusted Platform Module (TPM) Chip":
                system.mobo.tpm = d

ram_node = sub_nodes.find("MEMORY")
for r in ram_node:
    if r.tag == "Property":
        d = r.find("Description").text
        match r.find("Entry").text:
            case "Total Memory Size":
                system.ram.total = d

    if r.tag == "SubNode":
        ram = RAMStick()
        for s in r:
            if s.tag == "Property":
                d = s.find("Description").text
                match s.find("Entry").text:
                    case "Module Number":
                        ram.mod_num = d
                    case "Module Part Number":
                        ram.pn = d
                    case "Module Manufacturer":
                        ram.brand = d
                    case "Memory Type":
                        ram.mem_type = d
                    case "Memory Speed":
                        ram.speed = d
                    case "Module Size":
                        ram.size = d
        system.ram.sticks.append(ram)

cpu_node = sub_nodes.find("CPU")
for c in cpu_node:
    if c.tag == "Property":
        d = c.find("Description").text
        match c.find("Entry").text:
            case "Number Of Processor Cores":
                system.cpu.cores = d
            case "Number Of Logical Processors":
                system.cpu.lcores = d

    if c.tag == "SubNode":
        for _c in c:
            if _c.tag == "Property":
                d = _c.find("Description").text
                match _c.find("Entry").text:
                    case "Processor Name":
                        system.cpu.name = d
                    case "Original Processor Frequency [MHz]":
                        system.cpu.hz = f'{(int(d) / 1000)} Ghz'

gpu_node = sub_nodes.find("VIDEO")
for g in gpu_node:
    if g.tag == "SubNode":
        for _g in g:
            if _g.tag == "Property":
                d = _g.find("Description").text
                match _g.find("Entry").text:
                    case "Video Card":
                        system.gpu.name = d
                    case "Video Memory":
                        system.gpu.vram = d
                    case "Graphics Processor Clock":
                        system.gpu.clock = d
                    case "Graphics Memory Clock":
                        system.gpu.m_clock = d

drives_node = sub_nodes.find("DRIVES")

nvme_drives = drives_node[3]
for n_drive in nvme_drives:
    if n_drive.tag == "SubNode":
        drive = Drive()
        for p in n_drive:
            if p.tag == "Property":
                d = p.find("Description").text
                match p.find("Entry").text:
                    case "Drive Model":
                        drive.model = d
                    case "Drive Serial Number":
                        drive.sn = d
                    case "Drive Capacity":
                        drive.capacity = d
                    case "Drive Letter(s)":
                        for l in d.split(","):
                            part = Partition()
                            part.capacity = getSize(psutil.disk_usage(l).total).get("str")
                            part.free = getSize(psutil.disk_usage(l).free).get("str")
                            part.used = getSize(psutil.disk_usage(l).used).get("str")
                            part.mount_point = l
                            proc = subprocess.Popen(args=f"manage-bde -status {l}", stdout=subprocess.PIPE)
                            out, err = proc.communicate()
                            for line in out.decode("utf-8").split("\n"):
                                if "Protection Status:" in line:
                                    part.bitlocker = line.split(":")[1].strip()
                            drive.paritions.append(part)
        system.drives.nvme.append(drive)

sata_drives = drives_node[2]
for s_drive in sata_drives:
    if s_drive.tag == "SubNode":
        drive = Drive()
        for p in s_drive:
            if p.tag == "Property":
                d = p.find("Description").text
                match p.find("Entry").text:
                    case "Drive Model":
                        drive.model = d
                    case "Drive Serial Number":
                        drive.sn = d
                    case "Drive Capacity":
                        drive.capacity = d
                    case "Drive Letter(s)":
                        for l in d.split(","):
                            part = Partition()
                            part.capacity = getSize(psutil.disk_usage(l).total).get("str")
                            part.free = getSize(psutil.disk_usage(l).free).get("str")
                            part.used = getSize(psutil.disk_usage(l).used).get("str")
                            part.mount_point = l
                            proc = subprocess.Popen(args=f"manage-bde -status {l}", stdout=subprocess.PIPE)
                            out, err = proc.communicate()
                            for line in out.decode("utf-8").split("\n"):
                                if "Protection Status:" in line:
                                    part.bitlocker = line.split(":")[1].strip()
                            drive.paritions.append(part)
        system.drives.sata.append(drive)

ram_sticks = []

for stick in system.ram.sticks:
    ram_sticks.append(
        f"""
            Module Number:      {stick.mod_num}
            Part Number:        {stick.pn}
            Brand:              {stick.brand}
            Memory Type:        {stick.mem_type}
            Speed:              {stick.speed}
            Size:               {stick.size}
        """
    )

n_drives = []
for n in system.drives.nvme:
    parts = []
    for p in n.paritions:
        parts.append(
            f"""
                Volume:         {p.mount_point}
                Capacity:       {p.capacity}
                Used:           {p.used}
                Free:           {p.free}
                Bitlocker       {p.bitlocker}
            """
        )
    n_drives.append(
        f"""
            Model:              {n.model}
            Serial Number:      {n.sn}
            Capacity:           {n.capacity}
            Volumes:
                {"".join(parts)}
        """
    )

s_drives = []
for s in system.drives.sata:
    parts = []
    for p in s.paritions:
        parts.append(
            f"""
                Volume:         {p.mount_point}
                Capacity:       {p.capacity}
                Used:           {p.used}
                Free:           {p.free}
                Bitlocker:      {p.bitlocker}
            """
        )
    s_drives.append(
        f"""
            Model:              {s.model}
            Serial Number:      {s.sn}
            Capacity:           {s.capacity}
            Volumes:
                {"".join(parts)}
        """
    )

result = f"""System:
    Name:                       {system.name}
    OS:                         {system.os}
    Brand:                      {system.brand}
    UEFI:                       {system.uefi}
    Secure Boot:                {system.sb}

    Motherboard:
        Model:                  {system.mobo.model}
        Chipset:                {system.mobo.chipset}
        USB Version:            {system.mobo.usb_version}
        TPM:                    {system.mobo.tpm}

    RAM:
        Total Size:             {system.ram.total}
        Installed Sticks:
            {"".join(ram_sticks)}
    CPU:
        Name:                   {system.cpu.name}
        Cores:                  {system.cpu.cores}
        Logical Cores:          {system.cpu.lcores}
        Clock Speed:            {system.cpu.hz}

    GPU:
        Name:                   {system.gpu.name}
        Memory:                 {system.gpu.vram}
        Clock:                  {system.gpu.clock}
        Memory Clock:           {system.gpu.m_clock}

    Storage:
        NVMe Drives:
                {"".join(n_drives)}
        SATA Drives:
                {"".join(s_drives)}
    """

print(result)
with open("out.txt", "w") as f:
    f.write(result)
input()