#!/usr/bin/env python
|
#
|
# Copyright 2018 The Android Open Source Project
|
#
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
# you may not use this file except in compliance with the License.
|
#
|
# You may obtain a copy of the License at
|
#
|
# http://www.apache.org/licenses/LICENSE-2.0
|
#
|
# Unless required by applicable law or agreed to in writing, software
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
# See the License for the specific language governing permissions and
|
# limitations under the License.
|
"""Unit tests for at_auth_unlock."""
|
|
import argparse
|
import filecmp
|
import os
|
import shutil
|
import subprocess
|
import unittest
|
|
from at_auth_unlock import *
|
from Crypto.PublicKey import RSA
|
from unittest.mock import patch
|
|
|
def dataPath(file):
|
return os.path.join(os.path.dirname(__file__), 'data', file)
|
|
|
DATA_FILE_PIK_CERTIFICATE = dataPath('atx_pik_certificate.bin')
|
DATA_FILE_PUK_CERTIFICATE = dataPath('atx_puk_certificate.bin')
|
DATA_FILE_PUK_KEY = dataPath('testkey_atx_puk.pem')
|
DATA_FILE_UNLOCK_CHALLENGE = dataPath('atx_unlock_challenge.bin')
|
DATA_FILE_UNLOCK_CREDENTIAL = dataPath('atx_unlock_credential.bin')
|
|
|
def createTempZip(contents):
|
tempzip = tempfile.NamedTemporaryFile()
|
with zipfile.ZipFile(tempzip, 'w') as zip:
|
for arcname in contents:
|
zip.write(contents[arcname], arcname)
|
return tempzip
|
|
|
def validUnlockCredsZip():
|
return createTempZip({
|
'pik_certificate_v1.bin': DATA_FILE_PIK_CERTIFICATE,
|
'puk_certificate_v1.bin': DATA_FILE_PUK_CERTIFICATE,
|
'puk_v1.pem': DATA_FILE_PUK_KEY
|
})
|
|
|
class UnlockCredentialsTest(unittest.TestCase):
|
|
def testFromValidZipArchive(self):
|
with validUnlockCredsZip() as zip:
|
creds = UnlockCredentials.from_credential_archive(zip)
|
self.assertIsNotNone(creds.intermediate_cert)
|
self.assertIsNotNone(creds.unlock_cert)
|
self.assertIsNotNone(creds.unlock_key)
|
|
def testFromInvalidZipArchive(self):
|
with self.assertRaises(zipfile.BadZipfile):
|
UnlockCredentials.from_credential_archive(DATA_FILE_PUK_KEY)
|
|
def testFromArchiveMissingPikCertificate(self):
|
with createTempZip({
|
'puk_certificate_v1.bin': DATA_FILE_PUK_CERTIFICATE,
|
'puk_v1.pem': DATA_FILE_PUK_KEY
|
}) as zip:
|
with self.assertRaises(ValueError):
|
UnlockCredentials.from_credential_archive(zip)
|
|
def testFromArchiveMissingPukCertificate(self):
|
with createTempZip({
|
'pik_certificate_v1.bin': DATA_FILE_PIK_CERTIFICATE,
|
'puk_v1.pem': DATA_FILE_PUK_KEY
|
}) as zip:
|
with self.assertRaises(ValueError):
|
UnlockCredentials.from_credential_archive(zip)
|
|
def testFromArchiveMissingPuk(self):
|
with createTempZip({
|
'pik_certificate_v1.bin': DATA_FILE_PIK_CERTIFICATE,
|
'puk_certificate_v1.bin': DATA_FILE_PUK_CERTIFICATE,
|
}) as zip:
|
with self.assertRaises(ValueError):
|
UnlockCredentials.from_credential_archive(zip)
|
|
def testFromArchiveMultiplePikCertificates(self):
|
with createTempZip({
|
'pik_certificate_v1.bin': DATA_FILE_PIK_CERTIFICATE,
|
'pik_certificate_v2.bin': DATA_FILE_PIK_CERTIFICATE,
|
'puk_certificate_v1.bin': DATA_FILE_PUK_CERTIFICATE,
|
'puk_v1.pem': DATA_FILE_PUK_KEY
|
}) as zip:
|
with self.assertRaises(ValueError):
|
UnlockCredentials.from_credential_archive(zip)
|
|
def testFromArchiveMultiplePukCertificates(self):
|
with createTempZip({
|
'pik_certificate_v1.bin': DATA_FILE_PIK_CERTIFICATE,
|
'puk_certificate_v1.bin': DATA_FILE_PUK_CERTIFICATE,
|
'puk_certificate_v2.bin': DATA_FILE_PUK_CERTIFICATE,
|
'puk_v1.pem': DATA_FILE_PUK_KEY
|
}) as zip:
|
with self.assertRaises(ValueError):
|
UnlockCredentials.from_credential_archive(zip)
|
|
def testFromArchiveMultiplePuks(self):
|
with createTempZip({
|
'pik_certificate_v1.bin': DATA_FILE_PIK_CERTIFICATE,
|
'puk_certificate_v1.bin': DATA_FILE_PUK_CERTIFICATE,
|
'puk_v1.pem': DATA_FILE_PUK_KEY,
|
'puk_v2.pem': DATA_FILE_PUK_KEY
|
}) as zip:
|
with self.assertRaises(ValueError):
|
UnlockCredentials.from_credential_archive(zip)
|
|
def testFromFiles(self):
|
creds = UnlockCredentials(
|
intermediate_cert_file=DATA_FILE_PIK_CERTIFICATE,
|
unlock_cert_file=DATA_FILE_PUK_CERTIFICATE,
|
unlock_key_file=DATA_FILE_PUK_KEY)
|
self.assertIsNotNone(creds.intermediate_cert)
|
self.assertIsNotNone(creds.unlock_cert)
|
self.assertIsNotNone(creds.unlock_key)
|
|
def testInvalidPuk(self):
|
with self.assertRaises(ValueError):
|
UnlockCredentials(
|
intermediate_cert_file=DATA_FILE_PIK_CERTIFICATE,
|
unlock_cert_file=DATA_FILE_PUK_CERTIFICATE,
|
unlock_key_file=DATA_FILE_PUK_CERTIFICATE)
|
|
def testPukNotPrivateKey(self):
|
tempdir = tempfile.mkdtemp()
|
try:
|
with open(DATA_FILE_PUK_KEY, 'rb') as f:
|
key = RSA.importKey(f.read())
|
pubkey = os.path.join(tempdir, 'pubkey.pub')
|
with open(pubkey, 'wb') as f:
|
f.write(key.publickey().exportKey())
|
with self.assertRaises(ValueError):
|
UnlockCredentials(
|
intermediate_cert_file=DATA_FILE_PIK_CERTIFICATE,
|
unlock_cert_file=DATA_FILE_PUK_CERTIFICATE,
|
unlock_key_file=pubkey)
|
finally:
|
shutil.rmtree(tempdir)
|
|
def testWrongSizeCerts(self):
|
pik_cert = DATA_FILE_PIK_CERTIFICATE
|
tempdir = tempfile.mkdtemp()
|
try:
|
# Copy a valid cert and truncate a single byte from the end to create a
|
# too-short cert.
|
shortfile = os.path.join(tempdir, 'shortfile.bin')
|
shutil.copy2(pik_cert, shortfile)
|
with open(shortfile, 'ab') as f:
|
f.seek(-1, os.SEEK_END)
|
f.truncate()
|
with self.assertRaises(ValueError):
|
creds = UnlockCredentials(
|
intermediate_cert_file=shortfile,
|
unlock_cert_file=DATA_FILE_PUK_CERTIFICATE,
|
unlock_key_file=DATA_FILE_PUK_KEY)
|
with self.assertRaises(ValueError):
|
creds = UnlockCredentials(
|
intermediate_cert_file=DATA_FILE_PIK_CERTIFICATE,
|
unlock_cert_file=shortfile,
|
unlock_key_file=DATA_FILE_PUK_KEY)
|
|
# Copy a valid cert and append an arbitrary byte on the end to create a
|
# too-long cert.
|
longfile = os.path.join(tempdir, 'longfile.bin')
|
shutil.copy2(pik_cert, longfile)
|
with open(longfile, 'ab') as f:
|
f.write(b'\0')
|
with self.assertRaises(ValueError):
|
creds = UnlockCredentials(
|
intermediate_cert_file=longfile,
|
unlock_cert_file=DATA_FILE_PUK_CERTIFICATE,
|
unlock_key_file=DATA_FILE_PUK_KEY)
|
with self.assertRaises(ValueError):
|
creds = UnlockCredentials(
|
intermediate_cert_file=DATA_FILE_PIK_CERTIFICATE,
|
unlock_cert_file=longfile,
|
unlock_key_file=DATA_FILE_PUK_KEY)
|
finally:
|
shutil.rmtree(tempdir)
|
|
|
def writeFullUnlockChallenge(out_file, product_id_hash=None):
|
"""Helper function to create a file with a full AvbAtxUnlockChallenge struct.
|
|
Arguments:
|
product_id_hash: [optional] 32 byte value to include in the challenge as the
|
SHA256 hash of the product ID. If not provided, will default to the
|
product ID hash from the subject of DATA_FILE_PUK_CERTIFICATE.
|
"""
|
if product_id_hash is None:
|
with open(DATA_FILE_PUK_CERTIFICATE, 'rb') as f:
|
product_id_hash = GetAtxCertificateSubject(f.read())
|
assert len(product_id_hash) == 32
|
|
with open(out_file, 'wb') as out:
|
out.write(struct.pack('<I', 1))
|
out.write(product_id_hash)
|
with open(DATA_FILE_UNLOCK_CHALLENGE, 'rb') as f:
|
out.write(f.read())
|
|
|
class MakeAtxUnlockCredentialTest(unittest.TestCase):
|
|
def testCredentialIsCorrect(self):
|
with validUnlockCredsZip() as zip:
|
creds = UnlockCredentials.from_credential_archive(zip)
|
|
tempdir = tempfile.mkdtemp()
|
try:
|
challenge_file = os.path.join(tempdir, 'challenge')
|
writeFullUnlockChallenge(challenge_file)
|
challenge = UnlockChallenge(challenge_file)
|
out_cred = os.path.join(tempdir, 'credential')
|
|
# Compare unlock credential generated by function with one generated
|
# using 'avbtool make_atx_unlock_credential', to check correctness.
|
MakeAtxUnlockCredential(creds, challenge, out_cred)
|
self.assertTrue(filecmp.cmp(out_cred, DATA_FILE_UNLOCK_CREDENTIAL))
|
finally:
|
shutil.rmtree(tempdir)
|
|
def testWrongChallengeSize(self):
|
with validUnlockCredsZip() as zip:
|
creds = UnlockCredentials.from_credential_archive(zip)
|
|
tempdir = tempfile.mkdtemp()
|
try:
|
out_cred = os.path.join(tempdir, 'credential')
|
|
# The bundled unlock challenge is just the 16 byte challenge, not the
|
# full AvbAtxUnlockChallenge like this expects.
|
with self.assertRaises(ValueError):
|
challenge = UnlockChallenge(DATA_FILE_UNLOCK_CHALLENGE)
|
MakeAtxUnlockCredential(creds, challenge, out_cred)
|
finally:
|
shutil.rmtree(tempdir)
|
|
|
def makeFastbootCommandFake(testcase,
|
expect_serial=None,
|
error_on_command_number=None,
|
product_id_hash=None,
|
stay_locked=False):
|
"""Construct a fake fastboot command handler, to be used with unitttest.mock.Mock.side_effect.
|
|
This can be used to create a callable that acts as a fake for a real device
|
responding to the fastboot commands involved in an authenticated unlock. The
|
returned callback is intended to be used with unittest.mock.Mock.side_effect.
|
There are a number of optional arguments here that can be used to customize
|
the behavior of the fake for a specific test.
|
|
Arguments:
|
testcase: unittest.TestCase object for the associated test
|
expect_serial: [optional] Expect (and assert) that the fastboot command
|
specifies a specific device serial to communicate with.
|
error_on_command_number: [optional] Return a fastboot error (non-zero exit
|
code) on the nth (0-based) command handled.
|
stay_locked: [optional] Make the fake report that the device is still locked
|
after an otherwise successful unlock attempt.
|
"""
|
|
def handler(args, *extraArgs, **kwargs):
|
if error_on_command_number is not None:
|
handler.command_counter += 1
|
if handler.command_counter - 1 == error_on_command_number:
|
raise subprocess.CalledProcessError(
|
returncode=1, cmd=args, output=b'Fake: ERROR')
|
|
testcase.assertEqual(args.pop(0), 'fastboot')
|
if expect_serial is not None:
|
# This is a bit fragile in that, in reality, fastboot allows '-s SERIAL'
|
# to not just be the first arguments, but it works for this use case.
|
testcase.assertEqual(args.pop(0), '-s')
|
testcase.assertEqual(args.pop(0), expect_serial)
|
|
if args[0:2] == ['oem', 'at-get-vboot-unlock-challenge']:
|
handler.challenge_staged = True
|
elif args[0] == 'get_staged':
|
if not handler.challenge_staged:
|
raise subprocess.CalledProcessError(
|
returncode=1, cmd=args, output=b'Fake: No data staged')
|
|
writeFullUnlockChallenge(args[1], product_id_hash=product_id_hash)
|
handler.challenge_staged = False
|
elif args[0] == 'stage':
|
handler.staged_file = args[1]
|
elif args[0:2] == ['oem', 'at-unlock-vboot']:
|
if handler.staged_file is None:
|
raise subprocess.CalledProcessError(
|
returncode=1, cmd=args, output=b'Fake: No unlock credential staged')
|
|
# Validate the unlock credential as if this were a test key locked device,
|
# which implies tests that want a successful unlock need to be set up to
|
# use DATA_FILE_PUK_KEY to sign the challenge. Credentials generated using
|
# other keys will be properly rejected.
|
if not filecmp.cmp(handler.staged_file, DATA_FILE_UNLOCK_CREDENTIAL):
|
raise subprocess.CalledProcessError(
|
returncode=1, cmd=args, output=b'Fake: Incorrect unlock credential')
|
|
handler.locked = True if stay_locked else False
|
elif args[0:2] == ['getvar', 'at-vboot-state']:
|
return b'avb-locked: ' + (b'1' if handler.locked else b'0')
|
return b'Fake: OK'
|
|
handler.command_counter = 0
|
handler.challenge_staged = False
|
handler.staged_file = None
|
handler.locked = True
|
return handler
|
|
|
class AuthenticatedUnlockTest(unittest.TestCase):
|
|
@patch('subprocess.check_output')
|
def testUnlockWithZipArchive(self, mock_subp_check_output):
|
with validUnlockCredsZip() as zip:
|
mock_subp_check_output.side_effect = makeFastbootCommandFake(self)
|
self.assertEqual(main([zip.name]), 0)
|
self.assertNotEqual(mock_subp_check_output.call_count, 0)
|
|
@patch('subprocess.check_output')
|
def testUnlockDeviceBySerial(self, mock_subp_check_output):
|
with validUnlockCredsZip() as zip:
|
SERIAL = 'abcde12345'
|
mock_subp_check_output.side_effect = makeFastbootCommandFake(
|
self, expect_serial=SERIAL)
|
self.assertEqual(main([zip.name, '-s', SERIAL]), 0)
|
self.assertNotEqual(mock_subp_check_output.call_count, 0)
|
|
@patch('subprocess.check_output')
|
def testUnlockWithIndividualFiles(self, mock_subp_check_output):
|
mock_subp_check_output.side_effect = makeFastbootCommandFake(self)
|
self.assertEqual(
|
main([
|
'--pik_cert', DATA_FILE_PIK_CERTIFICATE, '--puk_cert',
|
DATA_FILE_PUK_CERTIFICATE, '--puk', DATA_FILE_PUK_KEY
|
]), 0)
|
self.assertNotEqual(mock_subp_check_output.call_count, 0)
|
|
@patch('subprocess.check_output')
|
def testFastbootError(self, mock_subp_check_output):
|
"""Verify that errors are handled properly if fastboot commands error out."""
|
with validUnlockCredsZip() as zip:
|
for n in range(5):
|
mock_subp_check_output.reset_mock()
|
mock_subp_check_output.side_effect = makeFastbootCommandFake(
|
self, error_on_command_number=n)
|
self.assertNotEqual(main([zip.name]), 0)
|
self.assertNotEqual(mock_subp_check_output.call_count, 0)
|
|
@patch('subprocess.check_output')
|
def testDoesntActuallyUnlock(self, mock_subp_check_output):
|
"""Verify fails if fake set to not actually unlock."""
|
with validUnlockCredsZip() as zip:
|
mock_subp_check_output.side_effect = makeFastbootCommandFake(
|
self, stay_locked=True)
|
self.assertNotEqual(main([zip.name]), 0)
|
self.assertNotEqual(mock_subp_check_output.call_count, 0)
|
|
@patch('subprocess.check_output')
|
def testNoCredentialsMatchDeviceProductID(self, mock_subp_check_output):
|
"""Test two cases where fake responds with a challenge that has a product ID hash which doesn't match the credentials used."""
|
# Case 1: Change the product ID hash that the fake responds with.
|
with validUnlockCredsZip() as zip:
|
mock_subp_check_output.side_effect = makeFastbootCommandFake(
|
self, product_id_hash=b'\x00' * 32)
|
self.assertNotEqual(main([zip.name]), 0)
|
self.assertNotEqual(mock_subp_check_output.call_count, 0)
|
|
# Case 2: Use credentials with a different product ID.
|
with createTempZip({
|
'pik_certificate_v1.bin': DATA_FILE_PIK_CERTIFICATE,
|
# Note: PIK cert used as PUK cert so subject (i.e. product ID hash) is
|
# different
|
'puk_certificate_v1.bin': DATA_FILE_PIK_CERTIFICATE,
|
'puk_v1.pem': DATA_FILE_PUK_KEY
|
}) as zip:
|
mock_subp_check_output.side_effect = makeFastbootCommandFake(self)
|
self.assertNotEqual(main([zip.name]), 0)
|
self.assertNotEqual(mock_subp_check_output.call_count, 0)
|
|
@patch('subprocess.check_output')
|
def testMatchingCredentialSelectedFromZipArchives(self,
|
mock_subp_check_output):
|
"""Test correct credential based on product ID hash used if multiple provided directly through arguments."""
|
with validUnlockCredsZip() as correctCreds, createTempZip({
|
'pik_certificate_v1.bin': DATA_FILE_PIK_CERTIFICATE,
|
# Note: PIK cert used as PUK cert so subject (i.e. product ID hash)
|
# doesn't match
|
'puk_certificate_v1.bin': DATA_FILE_PIK_CERTIFICATE,
|
'puk_v1.pem': DATA_FILE_PUK_KEY
|
}) as wrongCreds:
|
mock_subp_check_output.side_effect = makeFastbootCommandFake(self)
|
self.assertEqual(main([wrongCreds.name, correctCreds.name]), 0)
|
self.assertNotEqual(mock_subp_check_output.call_count, 0)
|
|
@patch('subprocess.check_output')
|
def testMatchingCredentialSelectedFromDirectory(self, mock_subp_check_output):
|
"""Test correct credential based on product ID hash used if multiple provided indirectly through a directory argument."""
|
with validUnlockCredsZip() as correctCreds, createTempZip({
|
'pik_certificate_v1.bin': DATA_FILE_PIK_CERTIFICATE,
|
# Note: PIK cert used as PUK cert so subject (i.e. product ID hash)
|
# doesn't match
|
'puk_certificate_v1.bin': DATA_FILE_PIK_CERTIFICATE,
|
'puk_v1.pem': DATA_FILE_PUK_KEY
|
}) as wrongCreds:
|
tempdir = tempfile.mkdtemp()
|
try:
|
shutil.copy2(correctCreds.name, tempdir)
|
shutil.copy2(wrongCreds.name, tempdir)
|
|
mock_subp_check_output.side_effect = makeFastbootCommandFake(self)
|
self.assertEqual(main([tempdir]), 0)
|
self.assertNotEqual(mock_subp_check_output.call_count, 0)
|
finally:
|
shutil.rmtree(tempdir)
|
|
@patch('subprocess.check_output')
|
def testMatchingCredentialSelectedFromEither(self, mock_subp_check_output):
|
"""Test correct credential based on product ID hash used if arguments give some combination of file and directory arguments."""
|
with validUnlockCredsZip() as correctCreds, createTempZip({
|
'pik_certificate_v1.bin': DATA_FILE_PIK_CERTIFICATE,
|
# Note: PIK cert used as PUK cert so subject (i.e. product ID hash)
|
# doesn't match
|
'puk_certificate_v1.bin': DATA_FILE_PIK_CERTIFICATE,
|
'puk_v1.pem': DATA_FILE_PUK_KEY
|
}) as wrongCreds:
|
# Case 1: Correct creds in directory, wrong in file arg
|
tempdir = tempfile.mkdtemp()
|
try:
|
shutil.copy2(correctCreds.name, tempdir)
|
|
mock_subp_check_output.side_effect = makeFastbootCommandFake(self)
|
self.assertEqual(main([wrongCreds.name, tempdir]), 0)
|
self.assertNotEqual(mock_subp_check_output.call_count, 0)
|
|
finally:
|
shutil.rmtree(tempdir)
|
|
# Case 2: Correct creds in file arg, wrong in directory
|
tempdir = tempfile.mkdtemp()
|
try:
|
shutil.copy2(wrongCreds.name, tempdir)
|
|
mock_subp_check_output.side_effect = makeFastbootCommandFake(self)
|
self.assertEqual(main([tempdir, correctCreds.name]), 0)
|
self.assertNotEqual(mock_subp_check_output.call_count, 0)
|
|
# Case 2: Correct creds in file arg, wrong in directory
|
finally:
|
shutil.rmtree(tempdir)
|
|
@patch('argparse.ArgumentParser.error')
|
def testArgparseDirectoryWithNoCredentials(self, mock_parser_error):
|
"""Test """
|
tempdir = tempfile.mkdtemp()
|
try:
|
# Make sure random files are ignored.
|
with open(os.path.join(tempdir, 'so_random'), 'w') as f:
|
f.write("I'm a random file")
|
|
mock_parser_error.side_effect = ValueError('ArgumentParser.error')
|
with self.assertRaises(ValueError):
|
main([tempdir])
|
self.assertEqual(mock_parser_error.call_count, 1)
|
finally:
|
shutil.rmtree(tempdir)
|
|
@patch('argparse.ArgumentParser.error')
|
def testArgparseMutualExclusionArchiveAndFiles(self, mock_parser_error):
|
mock_parser_error.side_effect = ValueError('ArgumentParser.error')
|
with self.assertRaises(ValueError):
|
main(['dummy.zip', '--pik_cert', DATA_FILE_PIK_CERTIFICATE])
|
self.assertEqual(mock_parser_error.call_count, 1)
|
|
@patch('argparse.ArgumentParser.error')
|
def testArgparseMutualInclusionOfFileArgs(self, mock_parser_error):
|
mock_parser_error.side_effect = ValueError('ArgumentParser.error')
|
with self.assertRaises(ValueError):
|
main(['--pik_cert', 'pik_cert.bin', '--puk_cert', 'puk_cert.bin'])
|
self.assertEqual(mock_parser_error.call_count, 1)
|
|
mock_parser_error.reset_mock()
|
with self.assertRaises(ValueError):
|
main(['--pik_cert', 'pik_cert.bin', '--puk', 'puk.pem'])
|
self.assertEqual(mock_parser_error.call_count, 1)
|
|
mock_parser_error.reset_mock()
|
with self.assertRaises(ValueError):
|
main(['--puk_cert', 'puk_cert.bin', '--puk', 'puk.pem'])
|
self.assertEqual(mock_parser_error.call_count, 1)
|
|
@patch('argparse.ArgumentParser.error')
|
def testArgparseMissingBundleAndFiles(self, mock_parser_error):
|
mock_parser_error.side_effect = ValueError('ArgumentParser.error')
|
with self.assertRaises(ValueError):
|
main(['-s', '1234abcd'])
|
self.assertEqual(mock_parser_error.call_count, 1)
|
|
|
if __name__ == '__main__':
|
unittest.main(verbosity=3)
|
