Refactor main script and update unit tests

List of changes:

* Renamed functions and variables for improved clarity and consistency.
* Fixed a bug where the script saved the boot image with the active slot suffix even if the user selected slot 'b'.
* Enhanced error handling and user feedback with ANSI escape codes.
* Improved partition detection logic and  introduced delays to improve UX.
* Expanded unit tests and adapted script modifications.

Signed-off-by: Abhijeet <[email protected]>

Author: gitclone-url

scripts/boot_image_extractor.py

@@ -1,79 +1,80 @@
 #!/usr/bin/env python3
 
-"""A script to extract boot image from either single or dual slotted Android devices with root access."""
+"""A script to extract boot images from any Android device with root access."""
 
 import os
 import sys
+import time
 import pyfiglet
 import subprocess
 
-def print_banner(name):
+def print_banner(title):
     max_width = os.get_terminal_size().columns
-    banner = pyfiglet.figlet_format(name, font='small', width=max_width)
+    banner = pyfiglet.figlet_format(title, font='small', width=max_width)
     print(banner.center(max_width))
-    
-def exit_with_error(error, reason):
-    print("\nError:", error)
-    print("\nReason:", reason)
+
+def exit_with_error(error_message, error_detail):
+    print("\033[91m\nError:\033[0m", error_message)
+    print("\nReason:", error_detail)
     sys.exit(1)
-      
-def extract_boot_image_dual_slot(boot_a_path, boot_b_path):
-    active_slot = subprocess.getoutput('getprop ro.boot.slot_suffix')
-    print("\nIt is recommended to extract the boot image according to the current active slot, which is ({}).\n".format(active_slot))
+
+def extract_boot_image_for_ab_device(boot_a_path, boot_b_path):
+    active_slot_suffix = subprocess.getoutput('getprop ro.boot.slot_suffix')
+    print(f"\n- Current active slot: ({active_slot_suffix}).")
+    time.sleep(1)
 
     while True:
-        chosen_slot = input("Which boot slot image would you like to extract? (a/b): ").lower()
-        if chosen_slot == 'a':
+        selected_slot = input("- Which boot slot image would you like to extract? (a/b): ").lower()
+        if selected_slot == 'a':
             boot_image_path = boot_a_path
             break
-        elif chosen_slot == 'b':
+        elif selected_slot == 'b':
             boot_image_path = boot_b_path
             break
         else:
-            print("Invalid input. Please choose either 'a' or 'b'.\n")
-            continue
+            print("Invalid input. Please enter 'a' or 'b'.\n")
 
-    print("\nExtracting the boot image from {}...".format(boot_image_path))
+    print(f"\n- Extracting the boot image from {boot_image_path}...")
+    time.sleep(1)
     try:
-        subprocess.check_call(['dd', 'if={}'.format(boot_image_path), 'of=./boot{}.img'.format(active_slot)])
-        print("Boot image successfully extracted and saved in your {} directory.".format(os.path.basename(os.getcwd())))
+        subprocess.check_call(['dd', f'if={boot_image_path}', f'of=./boot_{selected_slot}.img'])
+        print(f"\033[92m\n- Boot image successfully extracted and saved as boot_{selected_slot}.img in the current directory.\033[0m")
     except subprocess.CalledProcessError:
-        exit_with_error("Failed to extract the boot image", "dd command failed")
+        exit_with_error("Extraction failed", "The dd command did not complete successfully.")
 
-def extract_boot_image_single_slot(boot_path):
-    print("\nExtracting the boot image from {}...".format(boot_path))
+def extract_boot_image_for_legacy_device(boot_path):
+    print(f"\n- Extracting the boot image from {boot_path}...")
+    time.sleep(1)
     try:
-        subprocess.check_call(['dd', 'if={}'.format(boot_path), 'of=./boot.img'])
-        print("Boot image successfully extracted and saved in your {} directory.".format(os.path.basename(os.getcwd())))
+        subprocess.check_call(['dd', f'if={boot_path}', 'of=./boot.img'])
+        time.sleep(1)
+        print("\033[92m\n- Boot image successfully extracted and saved as boot.img in the current directory.\033[0m")
     except subprocess.CalledProcessError:
-        exit_with_error("Failed to extract the boot image", "dd command failed")
+        exit_with_error("Extraction failed", "The dd command did not complete successfully.")
 
 def main():
     if os.geteuid() != 0:
         exit_with_error("Insufficient privileges", "This script requires root access. Please run as root or use sudo.")
 
     print_banner("Boot Image Extractor")
+    time.sleep(1)
 
-    boot_names = ['boot', 'boot_a', 'boot_b']
-    for name in boot_names:
-        path = subprocess.getoutput('find /dev/block -type l -name {} -print | head -n 1'.format(name))
-        if path:
-            print("{} = {}".format(name, path))
-            if name == 'boot_a':
-                boot_a_path = path
-            elif name == 'boot_b':
-                boot_b_path = path
-            else:
-                boot_path = path
-
-    if 'boot_a_path' in locals() and 'boot_b_path' in locals():
-        print("\nDevice has dual boot slots.")
-        extract_boot_image_dual_slot(boot_a_path, boot_b_path)
-    elif 'boot_path' in locals():
-        print("\nDevice has a single boot slot.")
-        extract_boot_image_single_slot(boot_path)
+    boot_partitions = {}
+    for partition in ["boot", "boot_a", "boot_b"]:
+        partition_path = subprocess.getoutput(f"find /dev/block -type b -o -type l -iname '{partition}' -print -quit 2>/dev/null")
+        if partition_path:
+            boot_partitions[partition] = os.path.realpath(partition_path)
+            
+    if 'boot_a' in boot_partitions and 'boot_b' in boot_partitions:
+        print("\n- A/B partition style detected!.")
+        time.sleep(1)
+        extract_boot_image_for_ab_device(boot_partitions['boot_a'], boot_partitions['boot_b'])
+    elif 'boot' in boot_partitions:
+        print("\n- Legacy(non-A/B) partition style detected!.")
+        time.sleep(1)
+        extract_boot_image_for_legacy_device(boot_partitions['boot'])
     else:
-        exit_with_error("No boot slots found", "unable to find the symlinked boot slot files")
+        exit_with_error("No boot partition found", "Unable to locate block device files.")
 
 if __name__ == '__main__':
     main()

setup.py

@@ -9,8 +9,16 @@
 
 setup(
     name='Boot-Image-Extractor',
+    description='A tool to extract boot images from Android devices with root access.',
+    author='Abhijeet',
+    url='https://github.com/gitclone-url/Boot-image-Extractor',
     scripts=['scripts/boot_image_extractor.py'],
     install_requires=[
         'pyfiglet',
     ],
+    classifiers=[
+        'Programming Language :: Python :: 3',
+        'License :: OSI Approved :: MIT License',
+        'Operating System :: Android', 
+    ],
 )

tests/__init__.py

@@ -0,0 +1,5 @@
+import os
+import sys
+
+# Add the parent directory to the system path
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

tests/test_boot_image_extractor.py

@@ -1,48 +1,75 @@
 import unittest
-import os
-import sys
 from unittest.mock import patch, MagicMock
 
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
-
+from scripts.boot_image_extractor import (
+    print_banner, exit_with_error, extract_boot_image_for_legacy_device,
+    extract_boot_image_for_ab_device, main
+)
 
 class TestBootImageExtractor(unittest.TestCase):
-    """
-    Provides unit tests for the Boot Image Extractor script to ensure its proper functionality
-    in various scenarios. Additional test cases can be added to improve test coverage.
-    """
-    
-    def test_print_banner(self):
-        from scripts.boot_image_extractor import print_banner
-        with patch('builtins.print') as mock_print:
-            print_banner("Test Banner")
-            mock_print.assert_called_once()
-            
+    """Provides unit tests for the Boot Image Extractor script to ensure its proper functionality
+    in various scenarios. Additional test cases can be added to improve test coverage."""
+
+    def infinite_side_effect(self, *values):
+        while True:
+            for value in values:
+                yield value
+
+    @patch('builtins.print')
+    def test_print_banner(self, mock_print):
+        print_banner("Test Banner")
+        mock_print.assert_called()
+
     @patch('builtins.print')
     @patch('subprocess.getoutput')
     @patch('os.geteuid', MagicMock(return_value=0))
-    def test_extract_boot_image_single_slot(self, mock_getoutput, mock_print):
+    def test_extract_boot_image_for_legacy_device(self, mock_getoutput, mock_print):
         mock_getoutput.return_value = '/dev/block/boot'
         with patch('subprocess.check_call', MagicMock()) as mock_check_call:
-            with patch('os.path.basename', MagicMock(return_value='test')):
-                with patch('subprocess.getoutput', MagicMock(return_value='a')):
-                    from scripts.boot_image_extractor import extract_boot_image_single_slot
-                    extract_boot_image_single_slot('boot_path')
-                    mock_check_call.assert_called_with(['dd', 'if=boot_path', 'of=./boot.img'])
-                    
+            extract_boot_image_for_legacy_device('/dev/block/boot')
+            mock_check_call.assert_called_with(['dd', 'if=/dev/block/boot', 'of=./boot.img'])
+
     @patch('builtins.print')
     @patch('subprocess.getoutput')
     @patch('os.geteuid', MagicMock(return_value=0))
-    def test_extract_boot_image_dual_slot(self, mock_getoutput, mock_print):
-        mock_getoutput.side_effect = ['/dev/block/boot_a', '/dev/block/boot_b', 'a']
-        with patch('builtins.input', return_value='a'):
-            with patch('subprocess.check_call', MagicMock()) as mock_check_call:
-                with patch('os.path.basename', MagicMock(return_value='test')):
-                    with patch('subprocess.getoutput', MagicMock(return_value='a')):
-                        from scripts.boot_image_extractor import extract_boot_image_dual_slot
-                        extract_boot_image_dual_slot('boot_a_path', 'boot_b_path')
-                        mock_check_call.assert_called_with(['dd', 'if=boot_a_path', 'of=./boota.img'])
-                        
+    @patch('builtins.input', return_value='a')
+    def test_extract_boot_image_for_ab_device(self, mock_input, mock_getoutput, mock_print):
+        mock_getoutput.side_effect = ['_a', 'a']
+        with patch('subprocess.check_call', MagicMock()) as mock_check_call:
+            extract_boot_image_for_ab_device('/dev/block/boot_a', '/dev/block/boot_b')
+            mock_check_call.assert_called_with(['dd', 'if=/dev/block/boot_a', 'of=./boot_a.img'])
+
+    @patch('os.geteuid', MagicMock(return_value=0))
+    @patch('subprocess.getoutput', return_value='')
+    @patch('scripts.boot_image_extractor.print_banner', MagicMock())
+    @patch('scripts.boot_image_extractor.exit_with_error')
+    def test_main_no_boot_partition_found(self, mock_exit_with_error, mock_getoutput):
+        main()
+        mock_exit_with_error.assert_called_with("No boot partition found", "Unable to locate block device files.")
+        
+    @patch('os.geteuid', MagicMock(return_value=0))
+    @patch('subprocess.getoutput')
+    @patch('builtins.print')
+    def test_main_legacy_partition_style(self, mock_print, mock_getoutput):
+        mock_getoutput.side_effect = ['/dev/block/boot', '', '']
+        with patch('subprocess.check_call', MagicMock()) as mock_check_call:
+            main()
+            mock_print.assert_any_call("\n- Legacy(non-A/B) partition style detected!.")
+            mock_check_call.assert_called_with(['dd', 'if=/dev/block/boot', 'of=./boot.img'])
+            
+    @patch('os.geteuid', MagicMock(return_value=0))
+    @patch('subprocess.getoutput')
+    @patch('builtins.print')
+    @patch('builtins.input', return_value='a')
+    def test_main_ab_partition_style(self, mock_input, mock_print, mock_getoutput):
+        # Set the side_effect to our infinite_side_effect function
+        mock_getoutput.side_effect = self.infinite_side_effect('_a', '/dev/block/boot_a', '/dev/block/boot_b')
+        with patch('subprocess.check_call', MagicMock()) as mock_check_call:
+            main()
+            mock_print.assert_any_call("\n- A/B partition style detected!.")
+            mock_check_call.assert_called_with(['dd', 'if=/dev/block/boot_a', 'of=./boot_a.img'])        
 
 if __name__ == '__main__':
     unittest.main()
+
+