From 1a9a9cceefb93e1b8c3f4c70678d09491572507e Mon Sep 17 00:00:00 2001
From: Joe Leong <joeleong@users.noreply.github.com>
Date: Wed, 24 Jul 2019 13:31:11 -0700
Subject: [PATCH] Apply @Halastra suggestions from code review of #144

Co-Authored-By: Halastra <turchy@gmail.com>
---
 adb/android_pubkey.py | 35 ++++++++++++++++++++++++++++++-----
 setup.py              |  1 -
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/adb/android_pubkey.py b/adb/android_pubkey.py
index c1dddf0..f514343 100644
--- a/adb/android_pubkey.py
+++ b/adb/android_pubkey.py
@@ -8,7 +8,7 @@
 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
@@ -28,7 +28,6 @@
 from __future__ import print_function
 
 import os
-import six
 import base64
 import socket
 import struct
@@ -40,6 +39,17 @@
 # Size of an RSA modulus such as an encrypted block or a signature.
 ANDROID_PUBKEY_MODULUS_SIZE = (2048 // 8)
 
+# Python representation of "struct RSAPublicKey":
+ANDROID_PUBKEY_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 an encoded RSA key.
 ANDROID_PUBKEY_ENCODED_SIZE = \
     (3 * 4 + 2 * ANDROID_PUBKEY_MODULUS_SIZE)
@@ -52,7 +62,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]
@@ -71,7 +81,12 @@ def decode_pubkey(public_key):
     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]
+
+    key_struct = struct.unpack(ANDROID_PUBKEY_STRUCT, binary_key_data)
+    modulus_size_words, n0inv, modulus_bytes, rr_bytes, exponent = key_struct
+    assert modulus_size_words == ANDROID_PUBKEY_MODULUS_SIZE_WORDS  # Verify modulus length
+    modulus = reversed(modulus_bytes)
+    rr = reversed(rr_bytes)
     print('modulus_size_words:', hex(modulus_size_words))
     print('n0inv:', hex(n0inv))
     print('modulus: ', end='')
@@ -113,7 +128,17 @@ def encode_pubkey(private_key_path):
     key_buffer += _to_bytes(rr, ANDROID_PUBKEY_MODULUS_SIZE, 'little')
 
     key_buffer += struct.pack('<L', key.e)
-    return key_buffer
+
+    n_bytes = _to_bytes(key.n, ANDROID_PUBKEY_MODULUS_SIZE, 'little')
+    rr_bytes = _to_bytes(rr, ANDROID_PUBKEY_MODULUS_SIZE, 'little') 
+    return struct.pack(
+        ANDROID_PUBKEY_STRUCT,
+        ANDROID_PUBKEY_MODULUS_SIZE_WORDS, 
+        n0inv,
+        n_bytes,
+        rr_bytes,
+        key.e
+    )
 
 
 def get_user_info():
diff --git a/setup.py b/setup.py
index 1fe8337..7e6d67e 100644
--- a/setup.py
+++ b/setup.py
@@ -56,7 +56,6 @@
     install_requires = [
         'libusb1>=1.0.16',
         'pycryptodome',
-        'six',
         rsa_signer_library
     ],