import serial
import logging
import re
import telnetlib
import os
import subprocess
import tarfile
import base64
import string
import random
from time import sleep
from serial.serialutil import Timeout
import StringIO
import tempfile
logging.basicConfig(level=logging.DEBUG)
log = logging.getLogger(__name__)
class Rooter(object):
def __init__(self, **params):
if type(params['port']) is str:
params['port']=serial.Serial(
port=params['port'],
baudrate=115200
)
self._port = params['port']
self._ssh_pubkey_data = params['ssh_pubkey_data']
self._has_jtag = params['has_jtag']
self._check_uboot = params['check_uboot']
self._cleanup_payload = params['cleanup_payload']
self._reboot_after = params['reboot_after']
self._uboot_only = params['uboot_only']
self._boot_only = params['boot_only']
self._jtag_hardware = params['jtag_hardware']
def run(self):
#TODO: clean this up
port = self._port
ssh_pubkey_data = self._ssh_pubkey_data
has_jtag = self._has_jtag
check_uboot = self._check_uboot
cleanup_payload = self._cleanup_payload
reboot_after = self._reboot_after
uboot_only = self._uboot_only
boot_only = self._boot_only
if check_uboot:
uboot_passwords={
"2010.09-R6" : "f4E9J",
"2010.09-R8" : "3BHf2",
"2010.09" : "toon"
}
log.info("Waiting for Toon to restart")
uboot_version=self.read_uboot_version()
log.info("Toon has U-Boot version {}".format(uboot_version))
if uboot_version in uboot_passwords:
log.info("Using password to log in")
self.access_uboot(uboot_passwords[uboot_version])
if uboot_only:
log.info("Your Toon is now waiting in u-boot. You can you use a serial console for manual control.")
else:
self.patch_uboot()
if boot_only:
log.info("Your Toon is now booting into a serial console")
else:
log.info("Waiting for boot up")
self.write_payload()
self.patch_toon()
log.info("Your Toon is now rooted. Please wait for it to boot up and try to log in using SSH")
return
elif has_jtag is False:
log.error("Unable to log in using password (need JTAG, but it's disabled)")
return
if has_jtag:
log.info("Loading new bootloader")
self.start_bootloader("assets/u-boot.bin")
port.reset_input_buffer()
self._has_jtag = False
self._check_uboot = True
self.run()
else:
log.error("Need JTAG when rooting without manual reset")
return
def read_uboot_version(self):
version_line_match = re.compile(r'^U-Boot ([^ ]+)')
while True:
line = self._port.readline().strip()
match = version_line_match.match(line)
if match:
return match.group(1)
def access_uboot(self, password):
log.info("Logging in to U-Boot")
self._port.write(password)
self._port.flush()
log.debug(self._port.read_until("U-Boot>"))
log.debug("Logged in to U-Boot")
def patch_uboot(self):
port = self._port
log.info("Patching U-Boot")
port.reset_input_buffer()
sleep(0.1)
port.write("printenv\n")
port.flush()
add_misc_match = re.compile(r'^addmisc=(.+)$')
add_misc_val = None
sleep(0.5)
lines = port.read_until("U-Boot>")
log.debug(lines)
for line in lines.split('\n'):
line = line.strip()
log.debug(line)
match = add_misc_match.match(line)
if match:
add_misc_val = match.group(1)
if add_misc_val is None:
log.error("Could not find value for addmisc environment variable")
return
cmd = "setenv addmisc " + re.sub(r'([\$;])',r'\\\1', add_misc_val + " init=/bin/sh")
port.write(cmd + "\n")
port.flush()
log.debug(port.read_until("U-Boot>"))
port.write("run boot_nand\n")
port.flush()
def create_payload_tar(self):
tar_path = tempfile.mkstemp(suffix=".tar.gz")[1]
ssh_key = self._ssh_pubkey_data
with tarfile.open(tar_path, "w:gz") as tar:
tar.add('payload/', arcname='payload')
ssh_key_str = StringIO.StringIO(ssh_key)
info = tarfile.TarInfo(name="payload/id_rsa.pub")
info.size=len(ssh_key)
tar.addfile(tarinfo=info, fileobj=StringIO.StringIO(ssh_key))
return tar_path
def write_payload(self):
port = self._port
tar_path = self.create_payload_tar()
log.debug(port.read_until("/ # "))
port.write("base64 -d | tar zxf -\n")
port.flush()
#(tarr, tarw) = os.pipe()
#tar = tarfile.open(mode='w|gz', fileobj=tarw)
#tar.add("payload/patch_toon.sh")
log.info("Transferring payload")
with open(tar_path, 'r') as f:
base64.encode(f, port)
os.remove(tar_path)
port.flush()
port.reset_input_buffer()
port.write("\x04")
port.flush()
def patch_toon(self):
(port, clean_up, reboot) = (
self._port, self._cleanup_payload, self._reboot_after)
log.info("Patching Toon")
log.debug(port.read_until("/ # "))
password = ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(8))
port.write("sh payload/patch_toon.sh \"{}\"\n".format(password))
try:
while True:
line = read_until(port, ["/ # ", "\n"])
if line == "/ # ":
break
if line.startswith(">>>"):
log.info(line.strip())
else:
log.debug(line.strip())
except:
log.exception("Script failed")
sleep(5)
if clean_up:
log.info("Cleaning up")
port.write("rm -r payload\n")
log.debug(port.read_until("/ # "))
if reboot:
log.info("Rebooting")
port.write("/etc/init.d/reboot\n")
def start_bootloader(self, bin_path):
log.info("Starting openocd")
proc = subprocess.Popen([
'openocd',
'-s', '/usr/share/openocd',
'-f', 'assets/adapters/{}.cfg'.format(self._jtag_hardware),
'-f', 'assets/boards/ed20.cfg'
], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
try:
wait = 10
log.info("Waiting for {} seconds".format(wait))
sleep(wait)
client = telnetlib.Telnet('localhost', 4444)
log.debug(client.read_until("> "))
log.info("Halting CPU")
client.write("soft_reset_halt\n")
log.debug(client.read_until("> "))
sleep(0.1)
client.write("reset halt\n")
log.debug(client.read_until("> "))
sleep(0.1)
log.info("Loading new image to RAM")
client.write("load_image {} 0xa1f00000\n".format(bin_path))
log.debug(client.read_until("> "))
sleep(0.1)
log.info("Starting up new image")
client.write("resume 0xa1f00000\n")
except:
try:
log.exception(proc.communicate()[0])
except:
pass
proc.terminate()
raise
proc.terminate()
def read_until(port, terminators=None, size=None):
"""\
Read until any of the termination sequences is found ('\n' by default), the size
is exceeded or until timeout occurs.
"""
if not terminators:
terminators = ['\n']
terms = map(lambda t: (t, len(t)), terminators)
line = bytearray()
timeout = Timeout(port._timeout)
while True:
c = port.read(1)
if c:
line += c
for (terminator, lenterm) in terms:
if line[-lenterm:] == terminator:
# break does not work here because it will only step out of for
return bytes(line)
if size is not None and len(line) >= size:
break
else:
break
if timeout.expired():
break
return bytes(line)
