Refactor android_pubkey.py struct usage

With help from @Halastra suggestions from code review of #144

- Define ANDROID_RSAPUBLICKEY_STRUCT
- Remove use of six package
- Fix some flake8 warnings

Co-Authored-By: Halastra <[email protected]>

Author: joeleong

adb/android_pubkey.py

@@ -8,16 +8,16 @@
 https://github.com/aosp-mirror/platform_system_core/blob/c55fab4a59cfa461857c6a61d8a0f1ae4591900c/libcrypto_utils/android_pubkey.c
 
 typedef struct RSAPublicKey {
-    // Modulus length. This must be ANDROID_PUBKEY_MODULUS_SIZE.
+    // Modulus length. This must be ANDROID_PUBKEY_MODULUS_SIZE_WORDS
     uint32_t modulus_size_words;
 
     // Precomputed montgomery parameter: -1 / n[0] mod 2^32
     uint32_t n0inv;
 
-    // RSA modulus as a little-endian array.
+    // RSA modulus as a little-endian array
     uint8_t modulus[ANDROID_PUBKEY_MODULUS_SIZE];
 
-    // Montgomery parameter R^2 as a little-endian array of little-endian words.
+    // Montgomery parameter R^2 as a little-endian array of little-endian words
     uint8_t rr[ANDROID_PUBKEY_MODULUS_SIZE];
 
     // RSA modulus: 3 or 65537
@@ -28,7 +28,6 @@
 from __future__ import print_function
 
 import os
-import six
 import base64
 import socket
 import struct
@@ -40,10 +39,16 @@
 # Size of an RSA modulus such as an encrypted block or a signature.
 ANDROID_PUBKEY_MODULUS_SIZE = (2048 // 8)
 
-# Size of an encoded RSA key.
-ANDROID_PUBKEY_ENCODED_SIZE = \
-    (3 * 4 + 2 * ANDROID_PUBKEY_MODULUS_SIZE)
-  # (3 * sizeof(uint32_t) + 2 * ANDROID_PUBKEY_MODULUS_SIZE)
+# Python representation of "struct RSAPublicKey":
+ANDROID_RSAPUBLICKEY_STRUCT = (
+    '<'                 # Little-endian
+    'L'                 # uint32_t modulus_size_words;
+    'L'                 # uint32_t n0inv;
+    '{modulus_size}s'   # uint8_t modulus[ANDROID_PUBKEY_MODULUS_SIZE];
+    '{modulus_size}s'   # uint8_t rr[ANDROID_PUBKEY_MODULUS_SIZE];
+    'L'                 # uint32_t exponent;
+    ).format(modulus_size=ANDROID_PUBKEY_MODULUS_SIZE)
+
 
 # Size of the RSA modulus in words.
 ANDROID_PUBKEY_MODULUS_SIZE_WORDS = (ANDROID_PUBKEY_MODULUS_SIZE // 4)
@@ -52,7 +57,7 @@
 def _to_bytes(n, length, endianess='big'):
     """partial python2 compatibility with int.to_bytes
     https://stackoverflow.com/a/20793663"""
-    if six.PY2:
+    if not hasattr(n, 'to_bytes'):
         h = '{:x}'.format(n)
         s = ('0' * (len(h) % 2) + h).zfill(length * 2).decode('hex')
         return s if endianess == 'big' else s[::-1]
@@ -62,16 +67,11 @@ def _to_bytes(n, length, endianess='big'):
 def decode_pubkey(public_key):
     """decodes a public RSA key stored in Android's custom binary format"""
     binary_key_data = base64.b64decode(public_key)
-    key_struct = struct.unpack(('<LL' +
-        'B' * ANDROID_PUBKEY_MODULUS_SIZE +
-        'B' * ANDROID_PUBKEY_MODULUS_SIZE +
-        'L'), binary_key_data)
-    modulus_size_words = key_struct[0]
-    n0inv = key_struct[1]
-    modulus = reversed(key_struct[2: 2 + ANDROID_PUBKEY_MODULUS_SIZE])
-    rr = reversed(key_struct[2 + ANDROID_PUBKEY_MODULUS_SIZE:
-      2 + 2 * ANDROID_PUBKEY_MODULUS_SIZE])
-    exponent = key_struct[-1]
+    modulus_size_words, n0inv, modulus_bytes, rr_bytes, exponent = \
+        struct.unpack(ANDROID_RSAPUBLICKEY_STRUCT, binary_key_data)
+    assert modulus_size_words == ANDROID_PUBKEY_MODULUS_SIZE_WORDS
+    modulus = reversed(modulus_bytes)
+    rr = reversed(rr_bytes)
     print('modulus_size_words:', hex(modulus_size_words))
     print('n0inv:', hex(n0inv))
     print('modulus: ', end='')
@@ -83,15 +83,13 @@ def decode_pubkey(public_key):
 
 def decode_pubkey_file(public_key_path):
     with open(public_key_path, 'rb') as fd:
-      decode_pubkey(fd.read())
+        decode_pubkey(fd.read())
 
 
 def encode_pubkey(private_key_path):
     """encodes a public RSA key into Android's custom binary format"""
     key = Crypto.PublicKey.RSA.import_key(private_key_path)
 
-    # Store the modulus size.
-    key_buffer = struct.pack('<L', ANDROID_PUBKEY_MODULUS_SIZE_WORDS)
     # Compute and store n0inv = -1 / N[0] mod 2^32.
     # BN_set_bit(r32, 32)
     r32 = 1 << 32
@@ -101,19 +99,21 @@ def encode_pubkey(private_key_path):
     n0inv = Crypto.Util.number.inverse(n0inv, r32)
     # BN_sub(n0inv, r32, n0inv)
     n0inv = r32 - n0inv
-    key_buffer += struct.pack('<L', n0inv)
 
-    # Store the modulus.
-    key_buffer += _to_bytes(key.n, ANDROID_PUBKEY_MODULUS_SIZE, 'little')
     # Compute and store rr = (2^(rsa_size)) ^ 2 mod N.
     # BN_set_bit(rr, ANDROID_PUBKEY_MODULUS_SIZE * 8)
     rr = 1 << (ANDROID_PUBKEY_MODULUS_SIZE * 8)
     # BN_mod_sqr(rr, rr, key->n, ctx)
     rr = (rr ** 2) % key.n
-    key_buffer += _to_bytes(rr, ANDROID_PUBKEY_MODULUS_SIZE, 'little')
 
-    key_buffer += struct.pack('<L', key.e)
-    return key_buffer
+    return struct.pack(
+        ANDROID_RSAPUBLICKEY_STRUCT,
+        ANDROID_PUBKEY_MODULUS_SIZE_WORDS,
+        n0inv,
+        _to_bytes(key.n, ANDROID_PUBKEY_MODULUS_SIZE, 'little'),
+        _to_bytes(rr, ANDROID_PUBKEY_MODULUS_SIZE, 'little'),
+        key.e
+    )
 
 
 def get_user_info():
@@ -135,7 +135,7 @@ def write_public_keyfile(private_key_path, public_key_path):
         private_key = private_key_file.read()
 
     public_key = encode_pubkey(private_key)
-    assert len(public_key) == ANDROID_PUBKEY_ENCODED_SIZE
+    assert len(public_key) == struct.calcsize(ANDROID_RSAPUBLICKEY_STRUCT)
     with open(public_key_path, 'wb') as public_key_file:
         public_key_file.write(base64.b64encode(public_key))
         public_key_file.write(get_user_info().encode())

setup.py

@@ -56,7 +56,6 @@
     install_requires = [
         'libusb1>=1.0.16',
         'pycryptodome',
-        'six',
         rsa_signer_library
     ],