From bc2ab3cdf740534939f07c8dbcbe0df9af200cf9 Mon Sep 17 00:00:00 2001 From: Omur Ozbahceliler Date: Mon, 22 Jul 2024 23:26:24 +0100 Subject: [PATCH] colour added, surface scan promt added --- disk_scanner.py | 116 +++++++++++++++++++++++++++++++++++------------- main.py | 19 ++++++++ 2 files changed, 105 insertions(+), 30 deletions(-) diff --git a/disk_scanner.py b/disk_scanner.py index dfbf34f..ee920cb 100644 --- a/disk_scanner.py +++ b/disk_scanner.py @@ -18,7 +18,7 @@ def time_operation(operation, fd, sector, size): def write_sector(fd, sector, size): os.lseek(fd.fileno(), sector * size, os.SEEK_SET) - os.write(fd.fileno(), b'\0' * size) # Writing zero bytes to the sector + os.write(fd.fileno(), b'\0' * size) def read_sector(fd, sector, size): @@ -26,8 +26,7 @@ def read_sector(fd, sector, size): os.read(fd.fileno(), size) -def scan_disk(disk_path, sector_size, update_queue, lock): - # Get disk size in sectors +def scan_disk(disk_path, sector_size, update_queue, lock, stop_event): try: result = subprocess.run(['blockdev', '--getsz', disk_path], capture_output=True, text=True, check=True) total_sectors = int(result.stdout.strip()) @@ -39,6 +38,11 @@ def scan_disk(disk_path, sector_size, update_queue, lock): try: with open(disk_path, 'rb+', buffering=0) as fd: for sector in range(total_sectors): + if stop_event.is_set(): + with lock: + update_queue[disk_path]['stopped'] = True + break + write_time = time_operation(write_sector, fd, sector, sector_size) read_time = time_operation(read_sector, fd, sector, sector_size) @@ -84,41 +88,85 @@ def scan_disk(disk_path, sector_size, update_queue, lock): update_queue[disk_path]['error'] = f"Failed to open disk: {e}" -def update_ui(stdscr, update_queue, lock, disks): - while True: - stdscr.clear() - height, width = stdscr.getmaxyx() +def update_ui(stdscr, update_queue, lock, disk_map, stop_events): + curses.curs_set(0) + stdscr.nodelay(True) + curses.echo() - cols = width // 30 - rows = (len(disks) + cols - 1) + curses.start_color() + curses.init_pair(1, curses.COLOR_WHITE, curses.COLOR_BLACK) + curses.init_pair(2, curses.COLOR_CYAN, curses.COLOR_BLACK) + curses.init_pair(3, curses.COLOR_YELLOW, curses.COLOR_BLACK) + curses.init_pair(4, curses.COLOR_YELLOW, curses.COLOR_BLACK) + curses.init_pair(5, curses.COLOR_BLUE, curses.COLOR_BLACK) + curses.init_pair(6, curses.COLOR_MAGENTA, curses.COLOR_BLACK) + curses.init_pair(7, curses.COLOR_RED, curses.COLOR_BLACK) + curses.init_pair(8, curses.COLOR_RED, curses.COLOR_BLACK) - for idx, disk in enumerate(disks): - col = idx % cols - row = idx // cols + height, width = stdscr.getmaxyx() + max_width_per_disk = 30 + max_disks_per_row = width // max_width_per_disk + input_buffer = "" - x = col * 30 + while True: + stdscr.clear() + num_disks = len(disk_map) + num_rows = (num_disks + max_disks_per_row - 1) // max_disks_per_row + max_rows_for_display = (height - 1) // 10 + num_rows = min(num_rows, max_rows_for_display) + + for idx, (disk_num, disk) in enumerate(disk_map.items()): + row = idx // max_disks_per_row + col = idx % max_disks_per_row y = row * 10 + x = col * max_width_per_disk if y + 8 >= height: stdscr.addstr(height - 1, 0, "Not all disks are displayed.") break if 'error' in update_queue[disk]: - stdscr.addstr(y, x, f"{disk}: {update_queue[disk]['error']}") + stdscr.addstr(y, x, f"Disk {disk_num}: {update_queue[disk]['error']}", curses.color_pair(7)) continue - stdscr.addstr(y, x, f"Disk: {disk}") + if 'stopped' in update_queue[disk]: + stdscr.addstr(y, x, f"Disk {disk_num}: Stopped", curses.color_pair(5)) + continue + + stdscr.addstr(y, x, f"Disk {disk_num}: {disk}") with lock: stats = update_queue[disk] - stdscr.addstr(y + 1, x, f"<5ms = {stats['<5ms']}") - stdscr.addstr(y + 2, x, f"<10ms = {stats['<10ms']}") - stdscr.addstr(y + 3, x, f"<20ms = {stats['<20ms']}") - stdscr.addstr(y + 4, x, f"<50ms = {stats['<50ms']}") - stdscr.addstr(y + 5, x, f"<150ms = {stats['<150ms']}") - stdscr.addstr(y + 6, x, f"<500ms = {stats['<500ms']}") - stdscr.addstr(y + 7, x, f">500ms = {stats['>500ms']}") - stdscr.addstr(y + 8, x, f"BAD = {stats['bad']}") - + stdscr.addstr(y + 1, x, f"<5ms = {stats['<5ms']}", curses.color_pair(1)) + stdscr.addstr(y + 2, x, f"<10ms = {stats['<10ms']}", curses.color_pair(2)) + stdscr.addstr(y + 3, x, f"<20ms = {stats['<20ms']}", curses.color_pair(3)) + stdscr.addstr(y + 4, x, f"<50ms = {stats['<50ms']}", curses.color_pair(4)) + stdscr.addstr(y + 5, x, f"<150ms = {stats['<150ms']}", curses.color_pair(5)) + stdscr.addstr(y + 6, x, f"<500ms = {stats['<500ms']}", curses.color_pair(6)) + stdscr.addstr(y + 7, x, f">500ms = {stats['>500ms']}", curses.color_pair(7)) + stdscr.addstr(y + 8, x, f"BAD = {stats['bad']}", curses.color_pair(8) | curses.A_BOLD) + + prompt_str = "Press 'q' to quit. Enter disk number to stop: " + stdscr.addstr(height - 1, 0, prompt_str) + stdscr.addstr(height - 1, len(prompt_str), input_buffer) stdscr.refresh() + + key = stdscr.getch() + + if key == ord('q'): + break + + if key == curses.KEY_BACKSPACE or key == 127: + input_buffer = input_buffer[:-1] + + elif key == curses.KEY_ENTER or key == 10: + if input_buffer.isdigit(): + disk_num = int(input_buffer) + if disk_num in stop_events: + stop_events[disk_num].set() + input_buffer = "" + + elif key != -1 and chr(key).isdigit(): + input_buffer += chr(key) + time.sleep(1) @@ -127,13 +175,21 @@ def scan_disks(disks): update_queue = defaultdict(lambda: defaultdict(int)) lock = threading.Lock() + disk_map = {i: disk for i, disk in enumerate(disks)} + stop_events = {i: threading.Event() for i in disk_map} + threads = [] - for disk in disks: - t = threading.Thread(target=scan_disk, args=(disk, sector_size, update_queue, lock)) + for i, disk in disk_map.items(): + t = threading.Thread(target=scan_disk, args=(disk, sector_size, update_queue, lock, stop_events[i])) t.start() threads.append(t) - curses.wrapper(update_ui, update_queue, lock, disks) - - for t in threads: - t.join() + try: + curses.wrapper(update_ui, update_queue, lock, disk_map, stop_events) + except KeyboardInterrupt: + pass + finally: + for event in stop_events.values(): + event.set() + for t in threads: + t.join() diff --git a/main.py b/main.py index 511d84e..2b779b1 100644 --- a/main.py +++ b/main.py @@ -1,6 +1,7 @@ import signal import sys +import disk_scanner import diskforge from colorama import Fore, init @@ -12,6 +13,17 @@ def signal_handler(sig, frame): sys.exit(0) +def ask_user(prompt): + while True: + response = input(prompt).strip().lower() + if response in ['y', 'yes']: + return True + elif response in ['n', 'no']: + return False + else: + print("Please enter 'yes' or 'no'.") + + if __name__ == "__main__": signal.signal(signal.SIGINT, signal_handler) @@ -32,3 +44,10 @@ def signal_handler(sig, frame): print(f"{Fore.BLUE}====================================") diskforge.set_labels(disks) print(f"{Fore.BLUE}====================================") + + if ask_user("Would you like to surface scan the disks? [If you need to remove disks please do it now] (yes/no): "): + disks = diskforge.identify_disks() + disk_scanner.scan_disks(disks) + else: + print(f"{Fore.GREEN}Exiting without surface scan.") + sys.exit(0)