import threading
import os
import traceback

from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend

erasing_methods = {
    "2 прохода": {
        "num_passes": 2,
        "data": [b"\x00", "random"]
    },
    "Случайные данные": {
        "num_passes": 1,
        "data": ["random"]
    },
    "Гутмана": {
        "num_passes": 35,
        "data": [
            "random", "random", "random", "random", b"\x55", b"\xAA", b"\x92\x49\x24", b"\x49\x24\x92", b"\x24\x92\x49", b"\x00", b"\x11", b"\x22", b"\x33", b"\x44", b"\x55", b"\x66", b"\x77", b"\x88", b"\x99", b"\xAA", b"\xBB", b"\xCC", b"\xDD", b"\xEE", b"\xFF ",  b"\x92\x49\x24", b"\x49\x24\x92", b"\x24\x92\x49", b"\x6D\xB6\xDB", b"\xB6\xDB\x6D", b"\xDB\x6D\xB6", "random", "random", "random", "random"
            ]
    },
    "DoD Short": {
        "num_passes": 3,
        "data": [b"\x00", b"\xff", "random"]
    },
    "DoD 5220.22M": {
        "num_passes": 7,
        "data": [b"\x00", b"\xff", "random", b"\x00", b"\xff", "random", "random"]
    }
}

win_types = {
    0: "Unspecified",
    3: "HDD",
    4: "SSD",
    5: "SCM"
}

def get_random_bytes(size):
    seed=os.urandom(32)
    backend = default_backend()
    cipher = Cipher(algorithms.AES(seed), modes.CTR(b'\0'*16), backend=backend)
    encryptor = cipher.encryptor()
    nulls=b'\0'*(size)

    return encryptor.update(nulls)

class Drive():

    index: int

    path: str

    disk_type: str

    name: str

    block_size: int = 512

    capacity: int

    serial_num: str

    def __init__(self, name, path, disk_type, block_size, capacity, dev_id, index):
        self.name = name
        if isinstance(disk_type, int):
            self.disk_type = win_types[disk_type]
        else:
            self.disk_type = disk_type
        self.path = path
        if block_size is not None:
            self.block_size = block_size
        self.capacity = capacity
        self.serial_num = dev_id
        self.index = index

    def erase(self, method: str, progressbar, progress_label, result) -> threading.Thread:
        t = threading.Thread(target=self.__erasing, args=(method, progressbar,progress_label,result))
        t.start()
        return t

    def __erasing(self, method, progressbar, progress_label, result):
        write_size = self.block_size*32
        try:
            with open(self.path, "wb") as drive:
                for i in range(erasing_methods[method]["num_passes"]):
                    drive.seek(0)
                    prev_percent = 0
                    len_write = 0
                    progressbar.set(prev_percent)
                    progress_label.configure(text=f"{int(prev_percent*100)} %, {i+1}")
                    while len_write < self.capacity:
                        if erasing_methods[method]["data"][i] == "random":
                            data = get_random_bytes(write_size)
                        else:
                            data = erasing_methods[method]["data"][i] * write_size
                        if len(data) > self.capacity - len_write:
                            print(self.capacity - len_write)
                            data = data[:self.capacity - len_write]
                        drive.write(data)
                        len_write += len(data)
                        cur_percent = len_write/self.capacity
                        if cur_percent - prev_percent > 0.01:
                            progressbar.set(cur_percent)
                            progress_label.configure(text=f"{int(cur_percent*100)} %, {i+1}")
                            prev_percent = cur_percent
                        del data
            result["result"] = True
            result["err"] = ""
        except Exception as e:
            print(traceback.format_exc())
            result["result"] = False
            result["err"] = e