109 lines
5.3 KiB
Python
109 lines
5.3 KiB
Python
# SPDX-License-Identifier: MPL-2.0
|
|
import hashlib
|
|
import importlib.util
|
|
import json
|
|
from pathlib import Path
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
TOOLS = Path(__file__).resolve().parents[1]/'tools'
|
|
spec = importlib.util.spec_from_file_location('recovery', TOOLS/'restore-sargo.py')
|
|
recovery = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(recovery)
|
|
|
|
|
|
class RecoveryTests(unittest.TestCase):
|
|
def setUp(self):
|
|
temp = tempfile.TemporaryDirectory(); self.addCleanup(temp.cleanup)
|
|
self.root = Path(temp.name)
|
|
self.image = self.root/'boot_a.img'
|
|
self.image.write_bytes(b'ANDROID!'+bytes(4088))
|
|
self.size = patch.object(recovery, 'BOOT_SIZE', 4096)
|
|
self.size.start(); self.addCleanup(self.size.stop)
|
|
self.manifest = {'schema': 1, 'device': 'sargo', 'partition': 'boot_a', 'slot': 'a',
|
|
'image': 'boot_a.img', 'bytes': 4096,
|
|
'serialSha256': hashlib.sha256(b'PHONE').hexdigest(),
|
|
'imageSha256': recovery.digest(self.image), 'backupManifestSha256': 'a'*64}
|
|
self.save()
|
|
|
|
def save(self):
|
|
(self.root/'recovery.json').write_text(json.dumps(self.manifest))
|
|
|
|
def test_local_review_does_not_require_checkout_or_phone(self):
|
|
with patch.object(sys, 'argv', ['restore', '--kit', str(self.root)]), \
|
|
patch.object(recovery.subprocess, 'run') as run, \
|
|
patch.object(recovery.subprocess, 'check_output') as output, \
|
|
patch.object(recovery.os, 'umask'):
|
|
recovery.main()
|
|
run.assert_not_called(); output.assert_not_called()
|
|
|
|
def test_corrupted_backup_rejected(self):
|
|
self.image.write_bytes(b'ANDROID!'+b'x'*4088)
|
|
with self.assertRaisesRegex(ValueError, 'checksum mismatch'):
|
|
recovery.checked_kit(self.root)
|
|
|
|
def test_incomplete_export_cannot_restore(self):
|
|
(self.root/'.incomplete').touch()
|
|
with patch.object(sys, 'argv', ['restore', '--kit', str(self.root), '--restore']), \
|
|
patch.object(recovery, 'checked_device') as device:
|
|
with self.assertRaisesRegex(ValueError, 'did not finish'):
|
|
recovery.main()
|
|
device.assert_not_called()
|
|
|
|
def test_missing_header_even_with_matching_hash_rejected(self):
|
|
self.image.write_bytes(bytes(4096))
|
|
self.manifest['imageSha256'] = recovery.digest(self.image); self.save()
|
|
with self.assertRaisesRegex(ValueError, 'boot header'):
|
|
recovery.checked_kit(self.root)
|
|
|
|
def test_other_slot_and_path_are_rejected(self):
|
|
for key, value in (('partition', 'userdata'), ('slot', 'b'), ('image', '../boot_a.img')):
|
|
original = self.manifest[key]
|
|
self.manifest[key] = value; self.save()
|
|
with self.assertRaises(ValueError):
|
|
recovery.checked_kit(self.root)
|
|
self.manifest[key] = original
|
|
|
|
def test_wrong_phone_never_gets_a_write(self):
|
|
with patch.object(recovery.subprocess, 'check_output', return_value='OTHER fastboot\n'), \
|
|
patch.object(recovery.subprocess, 'run') as run:
|
|
with self.assertRaisesRegex(ValueError, 'not uniquely present'):
|
|
recovery.checked_device(self.manifest)
|
|
run.assert_not_called()
|
|
|
|
def test_locked_slot_b_wrong_product_and_wrong_partition_size_rejected(self):
|
|
original = {'product': 'sargo', 'unlocked': 'yes', 'current-slot': 'a',
|
|
'partition-size:boot_a': '0x1000'}
|
|
for name, value in (('product', 'bonito'), ('unlocked', 'no'),
|
|
('current-slot', 'b'), ('partition-size:boot_a', '0x2000')):
|
|
values = dict(original, **{name: value})
|
|
with patch.object(recovery.subprocess, 'check_output', return_value='PHONE fastboot\n'), \
|
|
patch.object(recovery, 'getvar', side_effect=lambda serial, key: values[key]):
|
|
with self.assertRaises(ValueError):
|
|
recovery.checked_device(self.manifest)
|
|
|
|
def test_restore_only_writes_boot_a_without_reboot(self):
|
|
with patch.object(sys, 'argv', ['restore', '--kit', str(self.root), '--restore']), \
|
|
patch.object(recovery, 'checked_device', return_value='PHONE'), \
|
|
patch.object(recovery.subprocess, 'run') as run, \
|
|
patch.object(recovery.os, 'umask'):
|
|
recovery.main()
|
|
run.assert_called_once_with(['fastboot', '-s', 'PHONE', 'flash', 'boot_a',
|
|
str(self.image.resolve())], check=True, timeout=90)
|
|
self.assertTrue(json.loads((self.root/'restore-report.json').read_text())['completed'])
|
|
|
|
def test_flash_failure_is_recorded_as_incomplete(self):
|
|
with patch.object(sys, 'argv', ['restore', '--kit', str(self.root), '--restore']), \
|
|
patch.object(recovery, 'checked_device', return_value='PHONE'), \
|
|
patch.object(recovery.subprocess, 'run', side_effect=subprocess.CalledProcessError(1, 'fastboot')), \
|
|
patch.object(recovery.os, 'umask'):
|
|
with self.assertRaises(subprocess.CalledProcessError):
|
|
recovery.main()
|
|
self.assertFalse(json.loads((self.root/'restore-report.json').read_text())['completed'])
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|