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 <98699436+gitclone-url@users.noreply.github.com>

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()
+
+