meta/meta-mediatek-mt2731/recipes-devtools/flashtool/files/flashimage.py

#!/usr/bin/python

import os
import os.path
import sys
import platform
import argparse
import shlex
import subprocess
import time
import re

################################################################################
# check python version
if sys.hexversion >= 0x02070000:
    pass
else:
    print('Please install Python 2.7.x or newer to run this script')
    exit(1)

# override input() in python2
try:
    input = raw_input
except NameError:
    pass

################################################################################

verbose = False
for_nata = False
skipFlash = False

fbtool = 'fbtool.py'
fastboot = 'fastboot'

system = platform.system()
machine = platform.machine()
product_out = os.path.abspath('.')

if 'Windows' in system or 'CYGWIN' in system:
    fastboot += '.exe'
elif 'Linux' in system:
    if 'x86_64' in machine:
        fastboot += '-linux-x86_64'
    elif 'arm' in machine or 'aarch64' in machine:
        fastboot += '-linux-arm'
elif 'Darwin' in system:
    fastboot += '-darwin'

# Generate image list from procedure list


def getImageList(procs):
    imgs = []
    try:
        for p in procs:
            if p[0] == 'fastboot' and p[1] == 'flash' and p[3] not in imgs:
                imgs.append(p[3])
    except Exception as e:
        print(e)
    return imgs


def call(cmd):
    '''cmd: the command list for subprocess.call'''
    if verbose:
        print('call:', ' '.join(cmd))
    if for_nata is True:
        return subprocess.call(cmd, stdout=sys.stdout, stderr=subprocess.STDOUT, shell=True)
    else:
        return subprocess.call(formatArgsForSubprocess(cmd))

def check_output(cmd):
    '''cmd: the command list for subprocess.check_output'''
    if verbose:
        print('check_output:', ' '.join(cmd))
    return subprocess.check_output(formatArgsForSubprocess(cmd), stderr=subprocess.STDOUT)


def checkImage(filename, needreboot=True, _verbose=False):
    filepath = os.path.join(product_out, filename)
    if os.path.exists(filepath):
        if _verbose:
            print(filename.rjust(39, ' ') + ':' + ' PASS')
        return filepath
    if needreboot:
        print(filename.rjust(39, ' ') + ':' + ' FAIL')
        if args.nata:
            print('[AT] Image not exist!')
        else:
            call([fastboot, 'reboot'])
        exit(1)
    return None

# return 0 if success


def cmdRun(cmd, dryrun=False):
    ret = 1
    raw_cmd = []
    if cmd[0] == 'daWait':
        daWait()
        print('')
        return 0
    elif cmd[0] == 'fbWait':
        fbWait()
        print('')
        return 0
    elif cmd[0] == 'fastboot':  # processing fastboot commands
        cmd[0] = fastboot
        if cmd[1] == 'flash':
            # check if image path exits and valid
            filepath = checkImage(cmd[3], False, False)
            if filepath != None:
                if 'CYGWIN' in system:
                    p = subprocess.check_output(
                        'cygpath --absolute --mixed %s' % filepath)
                    if p:
                        filepath = p.strip()
                raw_cmd += [cmd[0], cmd[1], cmd[2], filepath]
        else:
            raw_cmd += cmd
    else:
        print('FAIL: Unknown command!')
        return -1
    if dryrun:
        print(' '.join(raw_cmd))
        ret = 0
    else:
        ret = call(raw_cmd)
        # Bypass fastboot continue
        if cmd[1] == 'continue':
            ret = 0
        if ret == 0 and cmd[0] == 'fastboot' and cmd[1] == 'reboot-bootloader' and not dryrun:
            fbWait()
    return ret


def daWait(secs=60):
    print('Waiting for DA mode')
    ret = None
    for i in range(secs):
        sys.stdout.write('.')
        sys.stdout.flush()
        time.sleep(1)
        ret = check_output('python %s ' % fbtool)
        if ret != None and len(ret) != 0:
            print('')
            print('datool - device detected: ')
            break
    if ret == None:
        print('No device detected. Please ensure that datool is running')
        exit(1)


def fbWait(secs=60):
    print('Waiting for fastboot mode')
    ret = None
    for i in range(secs):
        sys.stdout.write('.')
        sys.stdout.flush()
        time.sleep(1)
        ret = check_output('%s devices' % fastboot)
        if ret != None and len(ret) != 0:
            print('')
            print('Fastboot - device detected: %s' % (ret.split())[0])
            break
    if ret == None:
        print('No device detected. Please ensure that fastboot is running on the target device')
        exit(1)


def cmdReboot(toBootloader=True):
    if toBootloader:
        call([fastboot, 'reboot-bootloader'])
        fbWait()
    else:
        call([fastboot, 'reboot'])


def formatArgsForSubprocess(cmd):
    if not 'Windows' in system:
        return shlex.split(cmd)
    else:
        return cmd


if __name__ == '__main__':
    # parse args
    parser = argparse.ArgumentParser(
        description='''
Auto device flasher, Python 2.7.x required
''',
        formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument('partition', nargs='?', default='all',
                        help='partition to flash [default: all] , not include test partition')
    parser.add_argument('-d', '--dryrun', action='store_true', default=False,
                        help='dryrun for debug, no image would be flashed')
    parser.add_argument('-u', '--user', action='store_true', default=False,
                        help='Flash user data partition')
    parser.add_argument('-b', '--boot', action='store_true', default=False,
                        help='Flash boot partition')
    parser.add_argument('-t', '--test', action='store_true', default=False,
                        help='Flash test partition')
    parser.add_argument('-v', '--verbose', action='store_true', default=False,
                        help='print more information while flashing')
    parser.add_argument('-n', '--nata', action = 'store_true', default = False,
                        help = 'Flash image in nata')
    parser.add_argument('--toolsdir', default=None,
                        help='''\
The tools dir where to find fbtool and fastboot.
Path priority order:
    1. --toolsdir specified
    2. current directory
    3. $PATH
''')
    parser.add_argument('--productdir', default=None,
                        help='''\
The product out directory where to find images.
Path priority order:
    1. --productdir specified
    2. current directory
''')

    args = parser.parse_args()
    verbose = args.verbose
    if args.dryrun:
        verbose = True

    print('')
    parser.print_usage()
    print('')
    print(''.center(80, '*'))
    print(('Running flasher on ' + platform.platform()).center(80))
    print(''.center(80, '*'))
    print('')

    if args.nata:
        for_nata = True
        args.toolsdir = os.path.abspath(os.path.dirname(__file__))
        args.productdir = args.toolsdir
        print ('[AT] Flash Dir: %s' %args.productdir)

    try:
        from flashproc import getFlashProc
    except ImportError as e:
        print('ImportError:', e)
        print('')
        exit(1)
    try:
        from flashproc import getFlashUserProc
    except:
        getFlashUserProc = getFlashProc
    try:
        from flashproc import getFlashBootProc
    except:
        getFlashBootProc = getFlashProc
    try:
        from flashproc import getFlashTestProc
    except:
        getFlashTestProc = getFlashProc

    # check flash tools
    toolsdir = ''
    try:
        if args.toolsdir:
            toolsdir = os.path.abspath(args.toolsdir)
            fbtool = os.path.join(toolsdir, fbtool)
            fastboot = os.path.join(toolsdir, fastboot)
            if not os.path.exists(fbtool) or not os.path.exists(fastboot):
                raise Exception(str(toolsdir))
        else:
            toolsdir = os.path.abspath('.')
            if os.path.exists(os.path.join(toolsdir, fbtool)) and os.path.exists(os.path.join(toolsdir, fastboot)):
                fbtool = os.path.join(toolsdir, fbtool)
                fastboot = os.path.join(toolsdir, fastboot)
    except Exception as e:
        print('Can not find fbtool or fastboot in %s' % str(e))
        print('')
        exit(1)

    devProduct = 'DEFAULT'

    procs = getFlashProc(devProduct)
    if args.user:
        procs = getFlashUserProc(devProduct)
    if args.boot:
        procs = getFlashBootProc(devProduct)
    if args.test:
        procs = getFlashTestProc(devProduct)
    if procs:
        if verbose:
            print('Flash procedure'.center(80))
            print(''.center(80, '-'))
            for p in procs:
                print('fastboot', ' '.join(p))
            print('')
    else:
        print('Can not retrieve flash procedure according to product type of', devProduct)
        print('Exit !')
        exit(1)

    # check image path
    if args.productdir:
        # take user specific product directory
        product_out = os.path.abspath(args.productdir)
    else:
        # check current directory
        product_out = os.path.abspath('.')
        for img in getImageList(procs):
            if checkImage(img, False, False) is None:
                product_out = None
                break
        if product_out is None:
            product_out = os.getenv('PRODUCT_OUT', '.')
    product_out = os.path.abspath(product_out)

    print(''.center(80, '*'))
    print('* flash images under:'.ljust(79, ' ') + '*')
    print('*     ' + product_out.ljust(73, ' ') + '*')
    print(''.center(80, '*'))
    print('')

    # check images
    print('Checking image'.center(80))
    print(''.center(80, '-'))
    images = getImageList(procs)
    try:
        for img in images:
            checkImage(img, _verbose=True)
    except Exception as e:
        print(e)
        time.sleep(2)
        if args.nata:
            print ('[AT] Check image fail, stop to flash!')
        else:
            cmdReboot(False)
            input('Fail, press enter to exit: ')
        exit(1)
    # For NATA
    if args.nata:
        #POWER OFF
        print ('[AT] POWER PIN KEY:0, %s' %time.ctime())
        time.sleep(2)
        #PRESS DOWNLOAD KEY
        print ('[AT] DOWNLOAD KEY:1, %s' %time.ctime())
        time.sleep(2)
        #POWER ON
        print ('[AT] POWER PIN KEY:1, %s' %time.ctime())
        time.sleep(2)

    # align 4kb for modem related images
    os.system('python align_4kb.py \"%s\"' % checkImage('md1img.img', False, False))

    # flash images
    print('')
    print('Start flashing'.center(80))
    print(''.center(80, '-'))
    try:
        for proc in procs:
            if skipFlash == False:
                if 0 != cmdRun(proc, args.dryrun):
                    raise Exception('<<FAILED>> %s' % ' '.join(proc))
            else:
                if proc[0] == 'fastboot' and proc[1] == 'flash':
                    print('Skip flashing', proc[3])
    except Exception as e:
        print(e)
        if args.nata:
            #RELEASE DOWNLOAD KEY
            print ('[AT] DOWNLOAD KEY:0, %s' %time.ctime())
            time.sleep(2)
        time.sleep(2)
        cmdReboot(False)
        exit(1)

    if args.nata:
        #RELEASE DOWNLOAD KEY
        print ('[AT] DOWNLOAD KEY:0, %s' %time.ctime())
        time.sleep(2)
    time.sleep(2)
    cmdReboot(False)
    print('')
    if args.nata:
        print ('[AT] Flash Success')
    else:
        print('Success')