From ef50b75ae115ae8162fa8b0a7b8c42b1d2db894b Mon Sep 17 00:00:00 2001
From: Plato Mavropoulos <platomav@users.noreply.github.com>
Date: Mon, 17 Jun 2024 23:55:57 +0300
Subject: [PATCH] Toshiba BIOS COM Extractor v4.0

Improved Toshiba COM detection (pattern only, no file extension)
Improved input object handling to support both paths and bytes
Populated code type hints and applied few small improvements
---
 README.md              |  2 +-
 Toshiba_COM_Extract.py | 43 +++++++++++++++++++-----------------------
 2 files changed, 20 insertions(+), 25 deletions(-)

diff --git a/README.md b/README.md
index 512326e..ff9ee9a 100644
--- a/README.md
+++ b/README.md
@@ -405,7 +405,7 @@ Moreover, you must have the following 3rd party tool at the "external" project d
 
 #### **Description**
 
-Parses Toshiba BIOS COM images and extracts their raw or compressed SPI/BIOS/UEFI firmware component. This utility is basically an easy to use python wrapper around [ToshibaComExtractor by LongSoft](https://github.com/LongSoft/ToshibaComExtractor). The output comprises only a final firmware component which is directly usable by end users.
+Parses Toshiba BIOS COM images and extracts their raw or compressed SPI/BIOS/UEFI firmware component. This utility is effectively a python wrapper around [ToshibaComExtractor by LongSoft](https://github.com/LongSoft/ToshibaComExtractor). The output comprises only a final firmware component which is directly usable by end users.
 
 #### **Usage**
 
diff --git a/Toshiba_COM_Extract.py b/Toshiba_COM_Extract.py
index 9ff8915..203b6b2 100644
--- a/Toshiba_COM_Extract.py
+++ b/Toshiba_COM_Extract.py
@@ -11,52 +11,47 @@ import os
 import subprocess
 
 from common.externals import get_comextract_path
-from common.path_ops import make_dirs, path_stem, path_suffixes
+from common.path_ops import make_dirs, path_stem, safe_name
 from common.patterns import PAT_TOSHIBA_COM
 from common.system import printer
 from common.templates import BIOSUtility
 from common.text_ops import file_to_bytes
 
-TITLE = 'Toshiba BIOS COM Extractor v3.0'
+TITLE = 'Toshiba BIOS COM Extractor v4.0'
 
 
-def is_toshiba_com(in_file):
+def is_toshiba_com(input_object: str | bytes | bytearray) -> bool:
     """ Check if input is Toshiba BIOS COM image """
 
-    buffer = file_to_bytes(in_file)
-
-    is_ext = path_suffixes(in_file)[-1].upper() == '.COM' if os.path.isfile(in_file) else True
-
-    is_com = PAT_TOSHIBA_COM.search(buffer)
-
-    return is_ext and is_com
+    return bool(PAT_TOSHIBA_COM.search(file_to_bytes(input_object)))
 
 
-def toshiba_com_extract(input_file, extract_path, padding=0):
+def toshiba_com_extract(input_object: str | bytes | bytearray, extract_path: str, padding: int = 0) -> int:
     """ Parse & Extract Toshiba BIOS COM image """
 
-    if not os.path.isfile(input_file):
-        printer('Error: Could not find input file path!', padding)
-
-        return 1
-
     make_dirs(extract_path, delete=True)
 
-    output_name = path_stem(input_file)
+    if isinstance(input_object, str) and os.path.isfile(input_object):
+        input_path: str = input_object
+    else:
+        input_path = os.path.join(extract_path, 'toshiba_bios.com')
 
-    output_file = os.path.join(extract_path, f'{output_name}.bin')
+        with open(input_path, 'wb') as input_buffer:
+            input_buffer.write(file_to_bytes(input_object))
+
+    output_path: str = os.path.join(extract_path, f'{safe_name(path_stem(input_path))}_extracted.bin')
 
     try:
-        subprocess.run([get_comextract_path(), input_file, output_file], check=True, stdout=subprocess.DEVNULL)
+        subprocess.run([get_comextract_path(), input_path, output_path], check=True, stdout=subprocess.DEVNULL)
 
-        if not os.path.isfile(output_file):
-            raise ValueError('EXTRACT_FILE_MISSING')
+        if not os.path.isfile(output_path):
+            raise FileNotFoundError('EXTRACTED_FILE_MISSING')
     except Exception as error:  # pylint: disable=broad-except
-        printer(f'Error: ToshibaComExtractor could not extract file {input_file}: {error}!', padding)
+        printer(f'Error: ToshibaComExtractor could not extract {input_path}: {error}!', padding)
 
-        return 2
+        return 1
 
-    printer(f'Succesfull {output_name} extraction via ToshibaComExtractor!', padding)
+    printer('Succesfull extraction via ToshibaComExtractor!', padding)
 
     return 0