Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

Commit fbc4792

Browse files
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>
1 parent 43d44c3 commit fbc4792

File tree

4 files changed

+117
-76
lines changed

4 files changed

+117
-76
lines changed

‎scripts/boot_image_extractor.py‎

Lines changed: 45 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1,79 +1,80 @@
11
#!/usr/bin/env python3
22

3-
"""A script to extract boot image from either single or dual slotted Android devices with root access."""
3+
"""A script to extract boot images from any Android device with root access."""
44

55
import os
66
import sys
7+
import time
78
import pyfiglet
89
import subprocess
910

10-
def print_banner(name):
11+
def print_banner(title):
1112
max_width = os.get_terminal_size().columns
12-
banner = pyfiglet.figlet_format(name, font='small', width=max_width)
13+
banner = pyfiglet.figlet_format(title, font='small', width=max_width)
1314
print(banner.center(max_width))
14-
15-
def exit_with_error(error, reason):
16-
print("\nError:", error)
17-
print("\nReason:", reason)
15+
16+
def exit_with_error(error_message, error_detail):
17+
print("\033[91m\nError:033円[0m", error_message)
18+
print("\nReason:", error_detail)
1819
sys.exit(1)
19-
20-
def extract_boot_image_dual_slot(boot_a_path, boot_b_path):
21-
active_slot = subprocess.getoutput('getprop ro.boot.slot_suffix')
22-
print("\nIt is recommended to extract the boot image according to the current active slot, which is ({}).\n".format(active_slot))
20+
21+
def extract_boot_image_for_ab_device(boot_a_path, boot_b_path):
22+
active_slot_suffix = subprocess.getoutput('getprop ro.boot.slot_suffix')
23+
print(f"\n- Current active slot: ({active_slot_suffix}).")
24+
time.sleep(1)
2325

2426
while True:
25-
chosen_slot = input("Which boot slot image would you like to extract? (a/b): ").lower()
26-
if chosen_slot == 'a':
27+
selected_slot = input("- Which boot slot image would you like to extract? (a/b): ").lower()
28+
if selected_slot == 'a':
2729
boot_image_path = boot_a_path
2830
break
29-
elif chosen_slot == 'b':
31+
elif selected_slot == 'b':
3032
boot_image_path = boot_b_path
3133
break
3234
else:
33-
print("Invalid input. Please choose either 'a' or 'b'.\n")
34-
continue
35+
print("Invalid input. Please enter 'a' or 'b'.\n")
3536

36-
print("\nExtracting the boot image from {}...".format(boot_image_path))
37+
print(f"\n- Extracting the boot image from {boot_image_path}...")
38+
time.sleep(1)
3739
try:
38-
subprocess.check_call(['dd', 'if={}'.format(boot_image_path), 'of=./boot{}.img'.format(active_slot)])
39-
print("Boot image successfully extracted and saved in your {} directory.".format(os.path.basename(os.getcwd())))
40+
subprocess.check_call(['dd', f'if={boot_image_path}', f'of=./boot_{selected_slot}.img'])
41+
print(f"033円[92m\n- Boot image successfully extracted and saved as boot_{selected_slot}.img in the current directory.033円[0m")
4042
except subprocess.CalledProcessError:
41-
exit_with_error("Failed to extract the boot image", "dd command failed")
43+
exit_with_error("Extraction failed", "The dd command did not complete successfully.")
4244

43-
def extract_boot_image_single_slot(boot_path):
44-
print("\nExtracting the boot image from {}...".format(boot_path))
45+
def extract_boot_image_for_legacy_device(boot_path):
46+
print(f"\n- Extracting the boot image from {boot_path}...")
47+
time.sleep(1)
4548
try:
46-
subprocess.check_call(['dd', 'if={}'.format(boot_path), 'of=./boot.img'])
47-
print("Boot image successfully extracted and saved in your {} directory.".format(os.path.basename(os.getcwd())))
49+
subprocess.check_call(['dd', f'if={boot_path}', 'of=./boot.img'])
50+
time.sleep(1)
51+
print("033円[92m\n- Boot image successfully extracted and saved as boot.img in the current directory.033円[0m")
4852
except subprocess.CalledProcessError:
49-
exit_with_error("Failed to extract the boot image", "dd command failed")
53+
exit_with_error("Extraction failed", "The dd command did not complete successfully.")
5054

5155
def main():
5256
if os.geteuid() != 0:
5357
exit_with_error("Insufficient privileges", "This script requires root access. Please run as root or use sudo.")
5458

5559
print_banner("Boot Image Extractor")
60+
time.sleep(1)
5661

57-
boot_names = ['boot', 'boot_a', 'boot_b']
58-
for name in boot_names:
59-
path = subprocess.getoutput('find /dev/block -type l -name {} -print | head -n 1'.format(name))
60-
if path:
61-
print("{} = {}".format(name, path))
62-
if name == 'boot_a':
63-
boot_a_path = path
64-
elif name == 'boot_b':
65-
boot_b_path = path
66-
else:
67-
boot_path = path
68-
69-
if 'boot_a_path' in locals() and 'boot_b_path' in locals():
70-
print("\nDevice has dual boot slots.")
71-
extract_boot_image_dual_slot(boot_a_path, boot_b_path)
72-
elif 'boot_path' in locals():
73-
print("\nDevice has a single boot slot.")
74-
extract_boot_image_single_slot(boot_path)
62+
boot_partitions = {}
63+
for partition in ["boot", "boot_a", "boot_b"]:
64+
partition_path = subprocess.getoutput(f"find /dev/block -type b -o -type l -iname '{partition}' -print -quit 2>/dev/null")
65+
if partition_path:
66+
boot_partitions[partition] = os.path.realpath(partition_path)
67+
68+
if 'boot_a' in boot_partitions and 'boot_b' in boot_partitions:
69+
print("\n- A/B partition style detected!.")
70+
time.sleep(1)
71+
extract_boot_image_for_ab_device(boot_partitions['boot_a'], boot_partitions['boot_b'])
72+
elif 'boot' in boot_partitions:
73+
print("\n- Legacy(non-A/B) partition style detected!.")
74+
time.sleep(1)
75+
extract_boot_image_for_legacy_device(boot_partitions['boot'])
7576
else:
76-
exit_with_error("No boot slots found", "unable to find the symlinked boot slot files")
77+
exit_with_error("No boot partition found", "Unable to locate block device files.")
7778

7879
if __name__ == '__main__':
7980
main()

‎setup.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,16 @@
99

1010
setup(
1111
name='Boot-Image-Extractor',
12+
description='A tool to extract boot images from Android devices with root access.',
13+
author='Abhijeet',
14+
url='https://github.com/gitclone-url/Boot-image-Extractor',
1215
scripts=['scripts/boot_image_extractor.py'],
1316
install_requires=[
1417
'pyfiglet',
1518
],
19+
classifiers=[
20+
'Programming Language :: Python :: 3',
21+
'License :: OSI Approved :: MIT License',
22+
'Operating System :: Android',
23+
],
1624
)

‎tests/__init__.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
import os
2+
import sys
3+
4+
# Add the parent directory to the system path
5+
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

‎tests/test_boot_image_extractor.py‎

Lines changed: 59 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,48 +1,75 @@
11
import unittest
2-
import os
3-
import sys
42
from unittest.mock import patch, MagicMock
53

6-
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
7-
4+
from scripts.boot_image_extractor import (
5+
print_banner, exit_with_error, extract_boot_image_for_legacy_device,
6+
extract_boot_image_for_ab_device, main
7+
)
88

99
class TestBootImageExtractor(unittest.TestCase):
10-
"""
11-
Provides unit tests for the Boot Image Extractor script to ensure its proper functionality
12-
in various scenarios. Additional test cases can be added to improve test coverage.
13-
"""
14-
15-
def test_print_banner(self):
16-
from scripts.boot_image_extractor import print_banner
17-
with patch('builtins.print') as mock_print:
18-
print_banner("Test Banner")
19-
mock_print.assert_called_once()
20-
10+
"""Provides unit tests for the Boot Image Extractor script to ensure its proper functionality
11+
in various scenarios. Additional test cases can be added to improve test coverage."""
12+
13+
def infinite_side_effect(self, *values):
14+
while True:
15+
for value in values:
16+
yield value
17+
18+
@patch('builtins.print')
19+
def test_print_banner(self, mock_print):
20+
print_banner("Test Banner")
21+
mock_print.assert_called()
22+
2123
@patch('builtins.print')
2224
@patch('subprocess.getoutput')
2325
@patch('os.geteuid', MagicMock(return_value=0))
24-
def test_extract_boot_image_single_slot(self, mock_getoutput, mock_print):
26+
def test_extract_boot_image_for_legacy_device(self, mock_getoutput, mock_print):
2527
mock_getoutput.return_value = '/dev/block/boot'
2628
with patch('subprocess.check_call', MagicMock()) as mock_check_call:
27-
with patch('os.path.basename', MagicMock(return_value='test')):
28-
with patch('subprocess.getoutput', MagicMock(return_value='a')):
29-
from scripts.boot_image_extractor import extract_boot_image_single_slot
30-
extract_boot_image_single_slot('boot_path')
31-
mock_check_call.assert_called_with(['dd', 'if=boot_path', 'of=./boot.img'])
32-
29+
extract_boot_image_for_legacy_device('/dev/block/boot')
30+
mock_check_call.assert_called_with(['dd', 'if=/dev/block/boot', 'of=./boot.img'])
31+
3332
@patch('builtins.print')
3433
@patch('subprocess.getoutput')
3534
@patch('os.geteuid', MagicMock(return_value=0))
36-
def test_extract_boot_image_dual_slot(self, mock_getoutput, mock_print):
37-
mock_getoutput.side_effect = ['/dev/block/boot_a', '/dev/block/boot_b', 'a']
38-
with patch('builtins.input', return_value='a'):
39-
with patch('subprocess.check_call', MagicMock()) as mock_check_call:
40-
with patch('os.path.basename', MagicMock(return_value='test')):
41-
with patch('subprocess.getoutput', MagicMock(return_value='a')):
42-
from scripts.boot_image_extractor import extract_boot_image_dual_slot
43-
extract_boot_image_dual_slot('boot_a_path', 'boot_b_path')
44-
mock_check_call.assert_called_with(['dd', 'if=boot_a_path', 'of=./boota.img'])
45-
35+
@patch('builtins.input', return_value='a')
36+
def test_extract_boot_image_for_ab_device(self, mock_input, mock_getoutput, mock_print):
37+
mock_getoutput.side_effect = ['_a', 'a']
38+
with patch('subprocess.check_call', MagicMock()) as mock_check_call:
39+
extract_boot_image_for_ab_device('/dev/block/boot_a', '/dev/block/boot_b')
40+
mock_check_call.assert_called_with(['dd', 'if=/dev/block/boot_a', 'of=./boot_a.img'])
41+
42+
@patch('os.geteuid', MagicMock(return_value=0))
43+
@patch('subprocess.getoutput', return_value='')
44+
@patch('scripts.boot_image_extractor.print_banner', MagicMock())
45+
@patch('scripts.boot_image_extractor.exit_with_error')
46+
def test_main_no_boot_partition_found(self, mock_exit_with_error, mock_getoutput):
47+
main()
48+
mock_exit_with_error.assert_called_with("No boot partition found", "Unable to locate block device files.")
49+
50+
@patch('os.geteuid', MagicMock(return_value=0))
51+
@patch('subprocess.getoutput')
52+
@patch('builtins.print')
53+
def test_main_legacy_partition_style(self, mock_print, mock_getoutput):
54+
mock_getoutput.side_effect = ['/dev/block/boot', '', '']
55+
with patch('subprocess.check_call', MagicMock()) as mock_check_call:
56+
main()
57+
mock_print.assert_any_call("\n- Legacy(non-A/B) partition style detected!.")
58+
mock_check_call.assert_called_with(['dd', 'if=/dev/block/boot', 'of=./boot.img'])
59+
60+
@patch('os.geteuid', MagicMock(return_value=0))
61+
@patch('subprocess.getoutput')
62+
@patch('builtins.print')
63+
@patch('builtins.input', return_value='a')
64+
def test_main_ab_partition_style(self, mock_input, mock_print, mock_getoutput):
65+
# Set the side_effect to our infinite_side_effect function
66+
mock_getoutput.side_effect = self.infinite_side_effect('_a', '/dev/block/boot_a', '/dev/block/boot_b')
67+
with patch('subprocess.check_call', MagicMock()) as mock_check_call:
68+
main()
69+
mock_print.assert_any_call("\n- A/B partition style detected!.")
70+
mock_check_call.assert_called_with(['dd', 'if=/dev/block/boot_a', 'of=./boot_a.img'])
4671

4772
if __name__ == '__main__':
4873
unittest.main()
74+
75+

0 commit comments

Comments
(0)

AltStyle によって変換されたページ (->オリジナル) /