path: root/images/add_adb_root/add_adb_root.py

#!/usr/bin/env python3
# Copyright (C) 2020 Denis 'GNUtoo' Carikli <GNUtoo@cyberdimension.org>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
import enum
import os
import re
import sh
import sys

# TODO:
# import tempfile

import files_checksums

class ImageType(enum.Enum):
    unknown = 0
    zImage = 1
    bootimage = 2

def usage(progname):
    print("{} <path/to/old/recovery.img> <path/to/new/recovery.img>".format(progname))
    sys.exit(1)

def automatically_identify_file(path):
    file_infos = {}

    # TODO: use file to get that
    file_infos['format'] = 'boot.img'

    # TODO: detect that and also allow 'boot'
    file_infos['type'] = 'recovery'

    # TODO: use file hash for that
    file_infos['release'] = None

    # TODO: is it really relevant?
    file_infos['target'] = None

    output = str(sh.unbootimg("-i", path)).split(os.linesep)

    for line in output:
        if line.startswith('kernel load addr:'):
            for word in line.split(" "):
                if word.startswith("0x"):
                    kernel_load_addr = int(word, 16)
                    # TODO: Add support for 64bit targets
                    file_infos['base_address'] = kernel_load_addr & 0xffff0000

        if line.startswith('cmdline:'):
            for word in line.split(" "):
                if word.startswith("`"):
                    file_infos['cmdline'] = \
                        word.replace("`", "").replace("'", "")

    return file_infos

def identify_file(path):
    output = sh.sha512sum(path).split(" ")
    checksum = output[0]

    # TODO: handle checksums not found
    file_infos = None
    try:
        file_infos = files_checksums.checksums[checksum]
    except:
        pass
    return file_infos

def identify_image_type(path):
    try:
        output = sh.unbootimg("-i", path)
    except Exception as e:
        if e.stderr == b'error: supplied file is not an Android boot image\n':
            return ImageType.zImage
        else:
            return ImageType.unknown

    return ImageType.bootimage

def check_file(file_path):
    file_infos = identify_file(file_path)
    if file_infos is None:
        file_infos = automatically_identify_file(file_path)

    if file_infos is None:
        print("/!\ TODO: Add support for that file by adding new checksums")
        sys.exit(1)

    if file_infos['type'] != 'recovery':
        print("/!\ The file is not a recovery")
        print("/!\ TODO: Add support new file types")
        sys.exit(1)

    return file_infos

def add_adb_to_ramdisk(ramdisk):
    sh.sed("s#ro.adb.secure=1#               #",
           "-i", ramdisk)

    sh.sed("s#ro.secure=1#ro.secure=0#",
           "-i", ramdisk)

    sh.sed("s#persist.sys.usb.config=none#persist.sys.usb.config=adb #",
           "-i", ramdisk)

def add_adb_to_zImage(input_file, output_file):
    tmpdir = str(sh.mktemp("-d")).replace(os.linesep, "")

    uncompressed_Image = None
    ramdisk = None

    args = ["-e", "-C", tmpdir, input_file]
    print(["binwalk"] + args)
    output = sh.binwalk(args)

    # Example: /tmp/tmp.dkbDvuu7PL/_recovery-i9100.img.extracted
    binwalk_dir = tmpdir + os.sep \
        + "_" + os.path.basename(input_file) + ".extracted"
    print("binwalk_dir: {}".format(binwalk_dir))

    files = os.listdir(binwalk_dir)
    for f in files:
        if f.endswith(".7z"):
            uncompressed_Image = binwalk_dir + os.sep + f[0:-3]

    # example: in
    # /tmp/tmp.dkbDvuu7PL/_recovery-i9100.img.extracted/1E74
    # we want 0x1E74 as int
    uncompressed_Image_offset = int("0x" + os.path.basename(uncompressed_Image),
                                    16)

    print("Uncompressed Image: {} @ {}".format(
        uncompressed_Image,
        str(hex(uncompressed_Image_offset))))

    # We want the ramdisk cpio file
    args = ["-e", "-C", binwalk_dir, uncompressed_Image]
    print(["binwalk"] + args)
    output = sh.binwalk(args)

    files = os.listdir(binwalk_dir)
    ramdisk_dir = binwalk_dir \
        + os.sep + "_" \
        + os.path.basename(uncompressed_Image) \
        + ".extracted"

    files = os.listdir(ramdisk_dir)
    for f in files:
        if f.endswith(".cpio"):
            ramdisk = ramdisk_dir + os.sep + f

    # example: in
    # /tmp/tmp.dkbDvuu7PL/_recovery-i9100.img.extracted/_1E74.extracted/32C9C.cpio
    # we want 0x32C9C as int
    ramdisk_offset = int("0x" + os.path.basename(ramdisk)[0:-5], 16)

    print("Ramdisk: {} @ {}".format(ramdisk, str(hex(ramdisk_offset))))

    add_adb_to_ramdisk(ramdisk)

    # ddrescue handles block size automatically
    # and doesn't truncate the file by default
    args = [ ramdisk,
             uncompressed_Image,
             "-o",
             str(ramdisk_offset)]
    print(["ddrescue"] + args)
    sh.ddrescue(args)

    # Issues:
    # no size_append like in scripts/Makefile.lib
    # Size too big
    sh.lzma("-9", uncompressed_Image)
    compressed_Image = uncompressed_Image + ".lzma"

    # TODO: Add the size of the file to the zImage

    if os.path.exists(output_file):
        sh.unlink(output_file)

    sh.cp(input_file, output_file)

    args = [compressed_Image,
            output_file ,
            "-o",
            str(uncompressed_Image_offset)]

    print(["ddrescue"] + args)

    sh.ddrescue(args)

def add_adb_to_bootimage(input_file, output_file):
    file_infos = check_file(input_file)

    # TODO:
    # tempfile.TemporaryDirectory().name
    tmpdir = str(sh.mktemp("-d")).replace(os.linesep, "")

    kernel = tmpdir + os.sep + "kernel.img"
    ramdisk = tmpdir + os.sep + "ramdisk.cpio"
    compressed_ramdisk = ramdisk + ".gz"

    # TODO: autodetect cmdline and base_address
    base_address = file_infos['base_address']
    cmdline = file_infos['cmdline']

    # TODO: check output
    output = sh.unbootimg("--kernel", kernel,
                          "--ramdisk", compressed_ramdisk,
                          "-i", input_file)

    sh.gunzip(compressed_ramdisk)

    # TODO:
    # file ramdisk.cpio

    add_adb_to_ramdisk(ramdisk)

    sh.gzip(ramdisk)

    sh.mkbootimg("--base", base_address,
                 "--kernel", kernel,
                 "--ramdisk", compressed_ramdisk,
                 "--cmdline={}".format(cmdline),
                 "-o", output_file)

if __name__ == "__main__":

    if len(sys.argv) != 3:
        usage(sys.argv[0])

    file_path = sys.argv[1]
    target_image = sys.argv[2]

    image_type = identify_image_type(file_path)
    if image_type == ImageType.zImage:
        add_adb_to_zImage(file_path, target_image)
    elif image_type == ImageType.bootimage:
        add_adb_to_bootimage(file_path, target_image)

    # TODO: clenup:
    # - remove tmpdir