ADB Backup Creator/Extractor

#!/usr/bin/python3

#credits
__title__       = "ADB Backup Creator/Extractor"
__author__      = "Visual Studio @ se7ensins.com"
__version__     = "1.0.0.0"
__description__ = "A script to create/extract ADB backups from Android devices"

#built-in imports
import zlib
from io import BytesIO
from os.path import isfile
from binascii import unhexlify
from argparse import ArgumentParser

#module imports
from Cryptodome.Cipher import AES
from Cryptodome.Util.Padding import unpad
from Cryptodome.Protocol.KDF import PBKDF2

#globals
KEY_BITS = 256
KEY_BYTES = int(KEY_BITS / 8)
PBKDF2_HASH_ROUNDS = 10000
BACKUP_FILE_HEADER_MAGIC = b"ANDROID BACKUP"

if __name__ == "__main__":
    #parse arguments
    parser = ArgumentParser(description="A script to decrypt/decompress android ADB backups")
    parser.add_argument("-e", "--create", action="store_true", help="Create a backup from scratch")
    parser.add_argument("-d", "--extract", action="store_true", help="Extract an existing backup")
    parser.add_argument("-i", "--in-file", type=str, help="The input file")
    parser.add_argument("-o", "--out-file", type=str, help="The output file")
    parser.add_argument("-k", "--key", type=str, help="The encryption key for the backup")
    args = parser.parse_args()

    #arg constraints
    assert args.create or args.extract, "No mode selected, use --encrypt or --decrypt!"
    assert isfile(args.in_file), "Input file doesn't exist!"
    assert not args.create, "Not implemented!"

    with open(args.in_file, "rb") as f:
        if args.extract:
            #read the header and make sure the magic exists
            assert f.readline().rstrip() == BACKUP_FILE_HEADER_MAGIC, "Invalid backup file magic!"

            #read out first half of the header
            backup_file_version = int(f.readline().rstrip())
            backup_file_compressed = bool(f.readline().rstrip())
            backup_file_encrypted = f.readline().rstrip() == b"AES-256"

            #the output
            backup_file_data = b""

            #if the file's encrypted then decrypt it
            if backup_file_encrypted:
                assert args.key, "No key provided and backup is encrypted, use --key!"

                #read out the rest of the header
                backup_file_user_password_salt = unhexlify(f.readline().rstrip())
                backup_file_master_key_checksum_salt = unhexlify(f.readline().rstrip())
                backup_file_pbkdf2_rounds = int(f.readline().rstrip())
                backup_file_user_password_iv = unhexlify(f.readline().rstrip())
                backup_file_enc_key_blob = unhexlify(f.readline().rstrip())
                backup_file_enc_contents = f.read()

                #derive the user password
                user_password = PBKDF2(args.key, backup_file_user_password_salt, KEY_BYTES, backup_file_pbkdf2_rounds)

                #decrypt the key blob
                backup_file_dec_key_blob = AES.new(user_password, AES.MODE_CBC, iv=backup_file_user_password_iv).decrypt(backup_file_enc_key_blob)
                with BytesIO(backup_file_dec_key_blob) as bio:
                    #read the key blob data
                    key_blob_iv = bio.read(bio.read(1)[0])
                    key_blob_master_key = bio.read(bio.read(1)[0])
                    key_blob_checksum = bio.read(bio.read(1)[0])

                    #check the master key's checksum against the stored one
                    #master_key_checksum = PBKDF2(key_blob_master_key, backup_file_master_key_checksum_salt, KEY_BYTES, PBKDF2_HASH_ROUNDS, hmac_hash_module=SHA1)

                    backup_file_data = unpad(AES.new(key_blob_master_key, AES.MODE_CBC, iv=key_blob_iv).decrypt(backup_file_enc_contents), AES.block_size)

            #if the file's compressed then decompress it
            if backup_file_compressed:
                #decompress the data
                backup_file_data = zlib.decompress(backup_file_data)

            #write the output file
            open(args.out_file, "wb").write(backup_file_data)