You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

603 lines
21 KiB
Python

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import sys
import types
import tarfile
import lzma
import ctypes
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
import gateway
from gateway import die
import read_info
import activate_boot
from devtree import *
import fdt
UIMAGE_MAGIC = b"\x27\x05\x19\x56"
FIT_MAGIC = FDT_MAGIC
HSQS_MAGIC = b"hsqs"
UBI_MAGIC = b"UBI#"
UBIv1_MAGIC = UBI_MAGIC + b"\x01\x00\x00\x00"
gw = gateway.Gateway()
if not gw.device_name:
die("Xiaomi Mi Wi-Fi device not found! (IP: {})".format(gw.ip_addr))
dn_dir = 'firmware/'
fn_local = None
dn_tmp = 'tmp/fw/'
fn_kernel = dn_tmp + 'kernel.bin'
fn_rootfs = dn_tmp + 'rootfs.bin'
fn_fitubi = dn_tmp + 'fitubi.bin'
os.makedirs(dn_dir, exist_ok = True)
os.makedirs(dn_tmp, exist_ok = True)
imglist = []
fn_list = [f for f in os.listdir(dn_dir) if os.path.isfile(os.path.join(dn_dir, f))]
for i, fname in enumerate(fn_list):
fname = dn_dir + fname
fsize = os.path.getsize(fname)
if fsize < 256*1024:
continue
with open(fname, "rb") as file:
data = file.read(256)
img = types.SimpleNamespace()
img.fn = fname
img.type = None
if data[:4] == b'HDR1' or data[:4] == b'HDR2':
img.type = 'stock'
if data[:10] == b"sysupgrade": # TAR
img.type = 'sysupgrade'
if data[:4] == UIMAGE_MAGIC: # uImage
img.type = 'factory'
if data[:4] == FIT_MAGIC: # factory squashfs image
img.type = 'factory'
if data[:8] == UBIv1_MAGIC: # rootfs_ubi
img.type = 'rootfs'
if img.type:
if len(imglist) == 0:
print('Image files in directory "{}":'.format(dn_dir))
print(' "{}"'.format(fname))
imglist.append(img)
if len(imglist) <= 0:
die('The firmware was not found in the "{}" folder!'.format(dn_dir))
c_stock = len([i for i in imglist if i.type == 'stock'])
c_sysupgrade = len([i for i in imglist if i.type == 'sysupgrade'])
c_factory = len([i for i in imglist if i.type == 'factory'])
c_rootfs = len([i for i in imglist if i.type == 'rootfs'])
if c_stock > 1 or c_factory > 1 or c_rootfs > 1 or c_sysupgrade > 1:
die('Too many different files in directory {}'.format(dn_dir))
if c_stock and c_sysupgrade and (c_factory or c_rootfs):
die('Too many different files in directory {}'.format(dn_dir))
if c_rootfs and c_factory == 0:
die('Kernel image not found! (only rootfs is present)')
dev = read_info.DevInfo(verbose = 0, infolevel = 1)
cpuarch = dev.info.cpu_arch
if cpuarch not in 'mips armv7 arm64':
die("Currently support only MIPS, ARMv7, ARM64 arch!")
class ImgHeader():
size = None # Image Data Size
os = None # OS type: IH_OS_LINUX = 5
arch = None # CPU arch: IH_CPU_MIPS = 5, IH_CPU_ARM = 2
type = None # Image Type: IH_TYPE_KERNEL = 2
comp = None # IH_COMP_NONE = 0, IH_COMP_GZIP = 1, IH_COMP_BZIP2 = 2, IH_COMP_LZMA = 3, IH_COMP_XZ = 5
class Image():
type = None # kernel or rootfs
ostype = None # 'stock', 'openwrt', 'padavan', 'pandorabox', etc
hdr = ImgHeader()
addr = None
addr2 = None # for kernel_stok/kernel_dup
fn_local = None
fn_remote = None
data = None
data2 = None # unpacked image
fit = False
dtb = None # device-tree
cmd = None
kernel = Image()
rootfs = Image()
if c_sysupgrade:
die("SysUpgrade firmware (TAR archives) are not supported!")
'''
fname = imglist[0].fn
file = tarfile.open(fname)
file.extractall(dn_tmp)
flist = file.getnames()
for fn in flist:
if fn.find('/') > 0:
bn = os.path.basename(fn)
if bn == 'kernel':
kernel.fn_local = dn_tmp + fn
if bn == 'root':
rootfs.fn_local = dn_tmp + fn
file.close()
if not kernel.fn_local:
die('Kernel image not found in TAR!')
with open(kernel.fn_local, "rb") as file:
kernel.data = file.read()
if not rootfs.fn_local:
die('Rootfs image not found in TAR!')
with open(rootfs.fn_local, "rb") as file:
rootfs.data = file.read()
# TODO: insert hsqs to UBI FS !!!
'''
def parse_fit(data, offset = 0):
if data[offset:offset+4] != FIT_MAGIC:
die('FIT: Incorrect image (0)')
hdrsize = ctypes.sizeof(fdt_header)
fit_hdr = fdt_header.from_buffer_copy(data[offset:offset+hdrsize])
fit_size = fit_hdr.totalsize
if fit_size < 1*1024*1024:
die('FIT: Incorrect image (1)')
if kernel.data:
die('FIT: Found second "kernel" section')
kernel.data = data[offset:offset + fit_size]
print('FIT size = 0x%X (%d KiB)' % (fit_size, fit_size // 1024))
fit_dt = fdt.parse_dtb(kernel.data)
#if fit_dt.root.nodes[0]._name != 'images':
# die('FIT: Incorrect image (4)')
fit_name = fit_dt.get_property('description').value
print('FIT name = "{}"'.format(fit_name))
kernel.fit = True
if 'OpenWrt FIT' in fit_name:
kernel.ostype = 'openwrt'
if not kernel.ostype:
die('FIT: Currently supported only OpenWrt FIT images!')
'''
fdt1 = dt.get_node('/images/fdt-1')
print('FDT desc = "{}"'.format(fdt1.get_property('description').value))
print('FDT type = "{}"'.format(fdt1.get_property('type').value))
print('FDT arch = "{}"'.format(fdt1.get_property('arch').value))
print('FDT compression = "{}"'.format(fdt1.get_property('compression').value))
if fdt1.get_property('type').value != 'flat_dt':
die('FIT: Incorrect image (6)')
kernel.hdr.arch = fdt1.get_property('arch').value
'''
if kernel.ostype == 'openwrt':
if len(data) - offset - fit_size < 1024:
return 1
rootfs_offset = data.find(UBIv1_MAGIC, offset + fit_size)
if rootfs_offset < 0:
return 1
rootfs_offset -= offset
print('FIT UBI offset = 0x%X' % rootfs_offset)
if rootfs.data:
die('FIT: Found two RootFS images')
rootfs.data = data[offset+rootfs_offset:]
rootfs.addr = rootfs_offset
return 2
return 1
def parse_factory(data, offset = 0):
if offset + 512 > len(data):
return -1
kernel_size = 0
if data[offset:offset+4] == UIMAGE_MAGIC: # uImage
pos = offset + 0x0C
kernel_size = int.from_bytes(data[pos:pos+4], byteorder='big')
kernel_size += 0x40
if kernel.data:
die('Found second "kernel" section!')
kernel.data = data[offset:offset+kernel_size]
if kernel_size > len(kernel.data):
die("Kernel header is incorrect!")
if data[offset:offset+4] == FIT_MAGIC: # factory squashfs image
return parse_fit(data, offset)
if kernel_size == 0:
die("Kernel header is incorrect!")
if kernel_size < 1*1024*1024:
die('Kernel image size is too small! (size: {} KB)'.format(kernel_size // 1024))
pos = 0x1C
kernel.hdr.size = kernel_size
kernel.hdr.arch = int.from_bytes(kernel.data[pos+1:pos+2], byteorder='little')
kernel.hdr.type = int.from_bytes(kernel.data[pos+2:pos+3], byteorder='little')
kernel.hdr.comp = int.from_bytes(kernel.data[pos+3:pos+4], byteorder='little')
kernel.hdr.name = kernel.data[0x20:0x40]
#print('cpu_arch = {}, img_type = {}, cmp_type = {}'.format(cpu_arch, img_type, cmp_type))
if kernel.hdr.type != 2: # IH_TYPE_KERNEL
die('Kernel type is incorrect!')
if kernel.hdr.arch == 5 and dev.info.cpu_arch != 'mips':
die('Kernel arch is not MIPS!')
if kernel.hdr.arch == 2 and dev.info.cpu_arch != 'armv7':
die('Kernel arch is not ARMv7!')
try:
iname = kernel.hdr.name.decode('ascii')
except Exception:
iname = None
if kernel.hdr.name[0:1] == b'\x03': # padavan kernel version
iname = None
if kernel.hdr.name[2:3] == b'\x03': # padavan fw version
ksize = int.from_bytes(kernel.hdr.name[0x1C:0x20], byteorder='big')
if ksize > 1*1024*1024 and ksize < kernel_size:
kernel.ostype = 'padavan'
try:
iname = kernel.hdr.name[4:0x1B].decode('ascii')
except Exception:
iname = None
kernel_size = ksize
#kernel.hdr.size = kernel_size
if kernel.data[ksize:ksize+4] != HSQS_MAGIC:
die('Incorrect padavan kernel image! RootFS not found!')
rootfs.data = kernel.data[ksize:]
if iname is None:
die('Incorrect kernel image name! (name: "{}")'.format(kernel.hdr.name))
print("Kernel image name = '{}'".format(iname))
kernel.hdr.name = iname
if kernel.ostype == 'padavan':
return 2
if c_stock:
#if kernel.hdr.comp == 0: # IH_COMP_NONE
# die("Stock Kernel must be packed!")
if iname.find('OpenWrt') < 0 or iname.find('Linux-3.') < 0:
die("Incorrect stock kernel image name!")
kernel.ostype = 'stock'
else:
if iname.find('PandoraBox') >= 0:
die('Pandora kernel not supported!')
if iname.find('OpenWrt') >= 0:
if iname.find('Linux-4.') > 0 or iname.find('Linux-5.') > 0 or iname.find('Linux-6.') > 0:
kernel.ostype = 'openwrt'
if not kernel.ostype:
die("Incorrect kernel image! Can't detect target OS type!")
if kernel_size > 4*1024*1024 - 128:
die('Kernel image size is too large! (size: {} KB)'.format(kernel_size // 1024))
if kernel_size > 0:
offset += kernel_size
data = data[offset:]
print('kernel_size = 0x%X' % kernel_size)
if len(data) < 1024:
return 1
offlist = []
offlist.append(data.find(UBIv1_MAGIC))
offlist.append(data.find(HSQS_MAGIC))
if len(offlist) == 0:
return 1
rootfs_offset = 0xFFFFFFFF
for i, off in enumerate(offlist):
if off >= 0 and off < rootfs_offset:
rootfs_offset = off
if rootfs_offset == 0xFFFFFFFF:
return 1
if rootfs.data:
die('Found two RootFS images!')
rootfs.data = data[rootfs_offset:]
return 2
return 0
if c_factory:
fname = [i for i in imglist if i.type == 'factory'][0].fn
with open(fname, "rb") as file:
data = file.read()
ret = parse_factory(data)
if ret < 1:
die('Kernel section not found!')
#print('c_factory OK')
if c_rootfs:
if rootfs.data:
die('Found two RootFS images!')
fname = [i for i in imglist if i.type == 'rootfs'][0].fn
with open(fname, "rb") as file:
data = file.read()
rootfs.data = data
if c_stock:
fname = imglist[0].fn
with open(fname, "rb") as file:
data = file.read()
imglst = []
for i in range(4):
p = 0x10 + i * 4
img = types.SimpleNamespace()
img.offset = int.from_bytes(data[p:p+4], byteorder='little')
if img.offset == 0:
continue
if img.offset > len(data) - 64:
die('Incorrect stock image! (2)')
if data[img.offset:img.offset+4] != b'\xBE\xBA\x00\x00':
die('Incorrect stock image! (3)')
img.size = int.from_bytes(data[img.offset+8:img.offset+8+4], byteorder='little')
hdr_size = 0x30
img.data = data[img.offset+hdr_size:img.offset+hdr_size+img.size]
if len(img.data) != img.size:
die('Incorrect stock image! (4)')
#print('offset = {} header = {}'.format("%08X" % (img.offset + hdr_size), img.data[:4]))
imglst.append(img)
if not imglst:
die('Incorrect stock image! (5)')
for i, img in enumerate(imglst):
if len(img.data) < 1*1024*1024: # skip uboot and other files
continue
if img.data[:4] == UIMAGE_MAGIC:
if kernel.data:
die('Incorrect stock image! (6)')
kernel.data = img.data
if img.data[:4] == UBI_MAGIC or img.data[:4] == HSQS_MAGIC:
if rootfs.data:
die('Incorrect stock image! (7)')
rootfs.data = img.data
if img.data[:4] == FDT_MAGIC:
if kernel.data or rootfs.data:
die('Incorrect stock image! (8)')
ret = parse_factory(img.data)
if ret < 1:
die('Kernel section not found!')
#print('c_stock OK')
if not kernel.data:
die('Kernel data not found!')
kernel.fn_remote = '/tmp/kernel.bin'
kernel.fn_local = fn_kernel
with open(kernel.fn_local, "wb") as file:
file.write(kernel.data)
if not rootfs.data:
die('RootFS data not found!')
rootfs.fn_remote = '/tmp/rootfs.bin'
rootfs.fn_local = fn_rootfs
with open(rootfs.fn_local, "wb") as file:
file.write(rootfs.data)
if kernel.data[:4] == FIT_MAGIC:
part_fw = dev.get_part('firmware')
part_fw1 = dev.get_part('firmware1')
if part_fw and part_fw1:
kernel.addr = part_fw['addr']
if not rootfs.addr:
die("FIT: Can't found addr for UBI image")
kernel_size = len(kernel.data)
kernel.data += (b"\x00" * (rootfs.addr - kernel_size)) + rootfs.data
rootfs.addr = None
kernel.fn_remote = '/tmp/fitubi.bin'
kernel.fn_local = fn_fitubi
with open(kernel.fn_local, "wb") as file:
file.write(kernel.data)
kernel.partname = "firmware" # kernel0
if not kernel.addr:
die("FIT: Can't found partition for flashing FIT image")
if kernel.data[:4] == UIMAGE_MAGIC:
data2 = kernel.data[0x40:]
img_comp = kernel.hdr.comp
if img_comp == 0:
loader_data = data2[:0x8000]
x1 = loader_data.find(b'Incorrect LZMA stream properties!') # b'OpenWrt kernel loader for MIPS based SoC'
x2 = loader_data.find(b'XZ-compressed data is corrupt')
if x1 < 0 and x2 < 0:
die('Kernel image not compressed! It is very strange!')
if x1 > 0 and x2 > 0:
die("Strange error! (1)")
if x1 > 0:
k1 = loader_data.find(b'\x00\x6D\x00\x00', x1) # LZMA prefix
k2 = loader_data.find(b'\x00\x5D\x00\x00', x1) # LZMA prefix
if k1 < 0 and k2 < 0:
die("Can't unpack kernel image! (lzma loader)")
if k1 > 0 and k2 > 0:
k = min(k1, k2)
else:
k = k1 if k1 > 0 else k2
data2 = data2[k+1:]
img_comp = 3 # IH_COMP_LZMA
if x2 > 0:
k = loader_data.find(b'\xFD7zXZ\x00', x2) # XZ prefix
if k < 0:
die("Can't unpack kernel image! (XZ loader)")
data2 = data2[k:]
img_comp = 5 # IH_COMP_XZ
if img_comp == 3: # IH_COMP_LZMA
kernel.data2 = lzma.decompress(data2)
if img_comp == 5: # IH_COMP_XZ
kernel.data2 = lzma.decompress(data2, lzma.FORMAT_XZ)
if kernel.data2:
with open(dn_tmp + 'kernel_unpacked.bin', "wb") as file:
file.write(kernel.data2)
if kernel.ostype == 'openwrt':
dtb = get_dtb(kernel.data, 32)
if not dtb and kernel.data2:
dtb = get_dtb(kernel.data2, 0)
if not dtb:
die("Can't found FDT (flattened device tree)")
kernel_part = get_dtb_part_info(dtb, "kernel")
if not kernel_part:
die('Cannot found "kernel" partition in DTB!')
print('part kernel = 0x%X (size: 0x%X)' % (kernel_part['addr'], kernel_part['size']))
kernel.addr = kernel_part['addr']
part = dev.get_part_by_addr(kernel.addr)
if not part:
die("Can't support flashing kernel to addr 0x%X" % kernel.addr)
kernel2_part = get_dtb_part_info(dtb, "kernel_dup")
if not kernel2_part:
kernel2_part = get_dtb_part_info(dtb, "kernel_stock")
if not kernel2_part:
die('Cannot found "kernel_dup"/"kernel_stock" partition in DTB!')
print('part kernel2 = 0x%X (size: 0x%X)' % (kernel2_part['addr'], kernel2_part['size']))
kernel.addr2 = kernel2_part['addr']
part = dev.get_part_by_addr(kernel.addr2)
if not part:
die("Can't support flashing kernel to addr 0x%X" % kernel.addr2)
ubi_part = get_dtb_part_info(dtb, "ubi")
if not ubi_part:
die('Cannot found "ubi" partition in DTB!')
print('part ubi = 0x%X (size: 0x%X)' % (ubi_part['addr'], ubi_part['size']))
rootfs.addr = ubi_part['addr']
part = dev.get_part_by_addr(rootfs.addr)
if not part:
die("Can't support flashing ubi to addr 0x%X" % rootfs.addr)
if len(rootfs.data) + 0x8000 >= part['size']:
die("Partition '%s' is too small (data size: 0x%X, part size: 0x%X)" % (part['name'], len(rootfs.data), part['size']))
if kernel.ostype == 'padavan':
kernel.addr = 0x600000
part = dev.get_part_by_addr(kernel.addr)
if not part:
die('Partition for addr {} not found'.format("0x%X" % kernel.addr))
part_size = part['size']
if len(kernel.data) > part_size:
rootfs.data = kernel.data[part_size:]
kernel.data = kernel.data[:part_size]
rootfs.addr = kernel.addr + part_size
part2 = dev.get_part_by_addr(rootfs.addr)
if not part2:
die('Partition for addr {} not found'.format("0x%X" % rootfs.addr))
else:
rootfs.addr = kernel.addr + part_size
part2 = dev.get_part_by_addr(rootfs.addr)
if not part2:
die('Partition for addr {} not found'.format("0x%X" % rootfs.addr))
part2_size = part2['size']
if part2_size > 8*1024*1024:
part2_size = 8*1024*1024
rootfs.data = b'\x00' * part2_size
with open(kernel.fn_local, "wb") as file:
file.write(kernel.data)
with open(rootfs.fn_local, "wb") as file:
file.write(rootfs.data)
dev.get_bootloader()
if not dev.bl.img:
die("Can't dump current bootloader!")
dev.get_env_list()
if not dev.env.fw.data or dev.env.fw.len <= 0:
die("Can't dump current NVRAM params!")
fw_num = None
if c_stock:
if dev.rootfs.num is None or dev.rootfs.num < 0:
die("Can't detect current booted rootfs!")
print("current flag_boot_rootfs = {}".format(dev.rootfs.num))
fw_num = 1 - dev.rootfs.num
#if dev.env.fw.var['flag_boot_rootfs'] == str(dev.rootfs.num):
# die("First, you should change the number of the boot kernel to {} !!!".format(fw_num))
kernel.partname = "kernel{}".format(fw_num)
kp = dev.get_part_num(kernel.partname)
if kp <= 0:
die("Partition {} not found!".format(kernel.partname))
kernel.addr = dev.partlist[kp]['addr']
kernel.cmd = 'mtd -e "{part}" write "{bin}" "{part}"'.format(part=kernel.partname, bin=kernel.fn_remote)
rootfs.partname = "rootfs{}".format(fw_num)
rp = dev.get_part_num(rootfs.partname)
if rp <= 0:
die("Partition {} not found!".format(rootfs.partname))
rootfs.addr = dev.partlist[rp]['addr']
rootfs.cmd = 'mtd -e "{part}" write "{bin}" "{part}"'.format(part=rootfs.partname, bin=rootfs.fn_remote)
if kernel.fit is True:
if not kernel.addr or not kernel.partname:
die('FIT: Unknown addr for flashing!')
fw_num = 0
kernel.cmd = 'mtd -e "{part}" write "{bin}" "{part}"'.format(part=kernel.partname, bin=kernel.fn_remote)
rootfs.cmd = None
elif kernel.ostype == 'openwrt' or kernel.ostype == 'padavan':
fw_addr = 0
if not kernel.addr or not rootfs.addr:
die('Unknown addr for flashing!')
part = dev.get_part_by_addr(kernel.addr)
if not part:
die('Partition for addr {} not found'.format("0x%X" % kernel.addr))
if part['size'] < len(kernel.data):
die('Partition size is too small!')
kernel.partname = part['name']
kernel.cmd = 'mtd -e "{part}" write "{bin}" "{part}"'.format(part=kernel.partname, bin=kernel.fn_remote)
if kernel.addr2:
part = dev.get_part_by_addr(kernel.addr2)
if not part:
die('Partition for addr {} not found'.format("0x%X" % kernel.addr2))
if part['size'] < len(kernel.data):
die('Partition size is too small!')
kernel.cmd += ' ; mtd -e "{part}" write "{bin}" "{part}"'.format(part=part['name'], bin=kernel.fn_remote)
part = dev.get_part_by_addr(rootfs.addr)
if not part:
die('Partition for addr {} not found'.format("0x%X" % rootfs.addr))
if part['size'] < len(rootfs.data):
die('Partition size is too small!')
rootfs.partname = part['name']
rootfs.cmd = 'mtd -e "{part}" write "{bin}" "{part}"'.format(part=rootfs.partname, bin=rootfs.fn_remote)
if dev.bl.type == 'breed':
fw_addr = None
if 'autoboot.command' in dev.env.breed.var:
cmd = dev.env.breed.var['autoboot.command']
lst = cmd.split(' ')
if len(lst) == 3:
try:
fw_addr = int(lst[2].strip(), 16)
except Exception:
fw_addr = None
if fw_addr and fw_addr == kernel.addr:
print("Breed boot address is correct! (addr: 0x%X)" % fw_addr)
else:
if fw_num is not None:
fw_addr = activate_boot.breed_boot_change(gw, dev, fw_num, None, None)
else:
fw_addr = activate_boot.breed_boot_change(gw, dev, None, kernel.addr, None)
pass
if not kernel.cmd or (not kernel.fit and not rootfs.cmd):
die("Flashing recipe unknown!")
gw.set_timeout(12)
gw.upload(kernel.fn_local, kernel.fn_remote)
if rootfs.cmd:
gw.upload(rootfs.fn_local, rootfs.fn_remote)
cmd = []
cmd.append("nvram set bootdelay=3")
cmd.append("nvram set boot_wait=on")
cmd.append("nvram set bootmenu_delay=30")
cmd.append("nvram set ssh_en=1")
cmd.append("nvram set uart_en=1")
cmd.append("nvram commit")
gw.run_cmd(cmd, timeout = 8)
if fw_num is not None:
print("Run scripts for change NVRAM params...")
activate_boot.uboot_boot_change(gw, fw_num)
print('Boot from partition "{}" activated.'.format(kernel.partname))
print('Writing kernel image to addr {} ...'.format("0x%08X" % kernel.addr))
print(" " + kernel.cmd)
gw.run_cmd(kernel.cmd, timeout = 34)
if rootfs.cmd:
print('Writing rootfs image to addr {} ...'.format("0x%08X" % rootfs.addr))
print(" " + rootfs.cmd)
gw.run_cmd(rootfs.cmd, timeout = 60)
print("The firmware has been successfully flashed!")
if kernel.fit:
print('Send command "reboot" via SSH ...')
gw.run_cmd("reboot -f")
print("Force REBOOT activated!")
else:
gw.run_cmd("sync ; umount -a", timeout = 5)
print("Please, reboot router!")