NexusOne刷机的一个脚本

#!/usr/bin/env python3
'''This program is used to flash an n1 image.'''

import os.path
import re
import time
import subprocess

# command to use
ADB = 'adb'
FASTBOOT = 'fastboot'
LOG_FILE = './flash.log'

class TimeoutError(Exception): pass

class Unbuffered:
    '''Unbuffered streams decorator.'''
    def __init__(self, stream):
        self.stream = stream

    def write(self, data):
        '''Each write followed a flush.'''
        self.stream.write(data)
        self.stream.flush()

    def __getattr__(self, attr):
        '''Keep other operations non-modified.'''
        return getattr(self.stream, attr)

def waitingDevice(mode='adb', timeout=60, logfile=None):
    '''Waiting nexus one to connect on the usb port.

    @param mode could be adb or fastboot.
    @param timeout.
    @logfile where to put log.
    @raises TimeoutError, ValueError.
    '''
    if mode == 'adb':
        cmd = ADB
        deviceLine = 1
        connectedIdentifer = 'device'
    elif mode == 'fastboot':
        cmd = FASTBOOT
        deviceLine = 0
        connectedIdentifer = 'fastboot'
    else:
        raise ValueError('unknown device connection mode: ' + mode)

    checkTime = 30
    checkInterval = timeout / checkTime
    i = 0
    while i < checkTime:
        lines = subprocess.Popen((cmd, 'devices'),
                                 stderr=logfile,
                                 stdout=subprocess.PIPE,
                                 universal_newlines=True).communicate()[0].splitlines()
        if len(lines) >= (deviceLine + 1):
            deviceInfo = re.split('\s+', lines[deviceLine])
            if (len(deviceInfo) >= 2 and
                deviceInfo[0].startswith('HT9CWP') and
                deviceInfo[1] == connectedIdentifer):
                return
        time.sleep(checkInterval)
        i += 1

    raise TimeoutError('device not connected in {} mode'.format(mode))

def reboot(mode='normal', logfile=None):
    '''Reboot the device.

    @param mode could be normal or fastboot.
    @param logfile where to put log.
    @raises TimeoutError, CalledProcessError, ValueError.'''

    if mode =='fastboot':
        waitingDevice('adb')
        subprocess.check_call((ADB, 'reboot', 'bootloader'),
                              stderr=subprocess.STDOUT,
                              stdout=logfile)
        waitingDevice('fastboot')
    elif mode == 'normal':
        waitingDevice('fastboot')
        subprocess.check_call((FASTBOOT, 'reboot'),
                              stderr=subprocess.STDOUT,
                              stdout=logfile)
        waitingDevice('adb')
    else:
        raise ValueError('unknown reboot mode: {}'.format(mode))

def eraseData(logfile=None):
    '''Erase device data.'''
    waitingDevice('fastboot')
    subprocess.check_call((FASTBOOT, 'erase', 'cache'),
                          stderr=subprocess.STDOUT,
                          stdout=logfile)
    subprocess.check_call((FASTBOOT, 'erase', 'userdata'),
                          stderr=subprocess.STDOUT,
                          stdout=logfile)

def flash(boot, system, userdata, logfile=None):
    '''Flash the specified image.'''
    waitingDevice('fastboot')
    subprocess.check_call((FASTBOOT, 'flash', 'boot', boot),
                          stderr=subprocess.STDOUT,
                          stdout=logfile)
    subprocess.check_call((FASTBOOT, 'flash', 'system', system),
                          stderr=subprocess.STDOUT,
                          stdout=logfile)
    subprocess.check_call((FASTBOOT, 'flash', 'userdata', userdata),
                          stderr=subprocess.STDOUT,
                          stdout=logfile)

def main(boot, system, userdata, logfile=None):
    '''Main entry point of flashing device.'''
    try:
        print('Waiting n1 to be connected ... ', end='')
        waitingDevice('adb', logfile=logfile)
        print('Done')

        print('Rebooting n1 to fastboot mode ... ', end='')
        reboot('fastboot',logfile=logfile)
        print('Done')

        print('Erasing cache and user data ... ', end='')
        eraseData(logfile=logfile)
        print('Done')

        print('Flashing device ... ', end='')
        flash(boot, system, userdata, logfile=logfile)
        print('Done')

        print('Rebooting n1 to normal mode ... ', end='')
        reboot('normal', logfile=logfile)
        print('Done')
    except Exception as e:
        print('\nFailed due to {}'.format(e))

def usage():
    '''Print usage information.'''
    print('Usage: {} --test/image_dir'.format(sys.argv[0]))

def test():
    try:
        '''Testing various functions of this module.'''
        print('\ttesting waitingDevice on adb mode ... ', end='')
        waitingDevice()
        print('ok')

        print('\ttesting reboot to fast boot mode ... ', end='')
        reboot('fastboot')
        print('ok')

        print('\ttesting reboot to normal mode ... ', end='')
        reboot('normal')
        print('ok')
    except Exception as e:
        print('\nFailed tests due to {}'.format(e))

if __name__ == '__main__':
    import sys
    # disable stdout bufferring
    sys.stdout = Unbuffered(sys.stdout)
    if len(sys.argv) != 2:
        usage()
        sys.exit(1)
    if sys.argv[1] == '--test':
        print("Running tests ...")
        test()
        sys.exit(0)

    # we will flash the device
    bootImg = os.path.join(sys.argv[1], "boot.img")
    systemImg = os.path.join(sys.argv[1], "system.img")
    userdataImg = os.path.join(sys.argv[1], "userdata.img")

    if (os.path.isfile(systemImg) and
        os.path.isfile(userdataImg) and
        os.path.isfile(bootImg)):
        print("Flashing image:\n\t'{}',\n\t'{}',\n\t'{}' to n1 ...".
              format(bootImg, systemImg, userdataImg))
        logfile = Unbuffered(open(LOG_FILE, 'a'))
        logfile.write("\n======= Begin of new log {} =======\n".
                      format(time.strftime("%a, %d %b %Y %H:%M:%S +0000", time.gmtime())))
        main(bootImg, systemImg, userdataImg, logfile)
        logfile.close()
        print("See {} for other details.".format(LOG_FILE))
    else:
        print("image file:\n\t'{}',\n\t'{}' or\n\t'{}' does not exists.".
              format(bootImg, systemImg, userdataImg))