#!/usr/bin/env python
'''
d6.py - Automatically fix d6 failures. 

Depends on:
    * Python 2.5 or greater
    * PySerial

You need a OLPC Serial Adapter or something compatible to connect to an XO-1.
See http://purl.org/lfaraone/olpc/serialinfo for more options.

This application is alpha. If you're on Ubuntu, you should remove the brltty 
application before continuing, as this has been known to cause compatibility
problems. 

'''
#
## Copyright (C) 2009 Luke Faraone
## Author: Luke Faraone <luke@laptop.org>
#
## This program is free software: you can redistribute it and/or modify
## it under the terms of the GNU 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
## General Public License for more details.
#
## You should have received a copy of the GNU General Public License
## along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

import time
import sys
import os
import traceback
import re

try:
    import serial
except ImportError: 
    print("Error: You need to have PySerial installed. For Ubuntu, run " + \
          "'sudo apt-get install python-serial'. " + \
	  " On an XO, run 'sudo yum install pyserial' " + \
	  " On Windows, visit http://pyserial.sf.net/ for more info. ")
    sys.exit(-1)

class XOLaptopSerialConnection():

    def __init__(self):
	self.ser = serial.Serial()
	self.ser.port = '/dev/ttyUSB0'
	self.ser.baudrate = 115200
	self.ser.timeout = 10

    def set_timeout(self,timeout):
	self.ser.timeout = timeout

    def set_port(self,port):
	self.ser.port = port

    def open(self):
	self.ser.open()
	return self.ser.isOpen()

    def is_open(self):
	return self.ser.isOpen()

    def close(self):
	self.ser.close()

    def do(self, action, output=None):
        self.write(action + '\n')
        self.wait_ok()

    def wait_ok(self, ok="ok"):
        #oldtimeout = device.getTimeout()
        #device.setTimeout(0)
        print "waiting on: " + ok
	# Make user the user can see this if we hang for some reason.
        sys.stdout.flush()
        while True:
	    line = self.ser.readline()
	    print "  -> %s" % line.strip()
	    if ok in line: break

    def disable_security(self):
        print "Attempting to disable security"
        self.write('disable-security\n')
        while True:
	    line = self.ser.readline()
	    print "  -> %s" % line.strip()
            if "No wp" in line:
                print "Security already disabled."
                return False
            elif "Restarting" in line:
                print "Security disabled."
                break
        return True

    def reboot(self):
        self.write('reboot\n')

    # A host machine can push commands down faster than OFW takes to process
    # the commands.  This causes input overflow in the serial input buffer
    # and we lose characters so add an intercharacter delay. Testing from the field
    # failed with a delay of 1ms so I've increased it 10ms.  Nothing in this is time
    # sensitive.
    # Normally when sending bulk data to OFW via serial port you would use download mode
    # 'dl' followed by a ctrl-d when the commands are finished.  Using this here for each
    # line seems like the wrong hammer. Plus you have to wait for OWF to be ready after
    # issueing 'dl'.
    def write(self, string):
	for each in string:
		self.ser.write(each)
		time.sleep(.01)
        
    def readline(self):
        string = self.ser.readline()
        return string

def unbrick(laptop):

	# regex to remove ansi codes.

	r= re.compile("\033\[[0-9;]+m")

	while True:
	    print("Connect a bricked XO and power it on.\n")
    
	    d6 = False
	    on = False
            laptop.set_timeout(10)
	    while True:
		line = laptop.readline()
		# OFW lines can contain ANSI sequnces that may mess up our terminal so remove them
		line = r.sub('',line)
		if line:
		    print "  -> %s" % line.strip()
		if "Page Fault" in line:
		    d6 = True
		    print "D6 bug detected"
		    break
		if "OpenFirmware" in line:
		    print "\nMachine booted.\n"
		    return
		if not line:
		    print "Read timeout. Type ctrl-C to exit."

	    if not d6:
		print "Your XO does not suffer from the D6 bug, or we missed the failure for some reason"
		return

	    if laptop.disable_security(): continue

	    print "selecting rtc"
	    laptop.do("select /rtc")
	    print "setting decimal"
	    laptop.do("decimal")
	    print "setting time"
	    laptop.do((time.strftime("%S %M %H %d %m %Y", time.gmtime()) + " set-time"))
	    laptop.do("time&date")
	    laptop.do(".date")
	    laptop.do(".time")
	    print "Date set. Rebooting"
	    laptop.reboot()
	    return

if __name__ == "__main__":

        VERSION = "0.1.3"
	# Try to use the first port thats not locked and that we can successfully 
	# open. This helps if you already have a serial connection in minicom
	# up to some thing else 

        print "D6 debrick Ver: " + VERSION
	ports = ['ttyUSB0','ttyUSB1','ttyUSB2']
	laptop = XOLaptopSerialConnection()

	try:
		for port in ports:
		    print "Trying port: %s" % port
		    devfile = '/dev/'+port
		    lckfile = '/var/lock/LCK..' + port
		    if not os.path.exists(devfile):
			print "%s not found" % port
			continue
		    if os.path.exists(lckfile):
			print "%s in use" % port
			continue
		    laptop.set_port(devfile)
		    if laptop.open():
			print "Using port: %s" % port
			break

		if not laptop.is_open():
			print "Can't find a usable serial port"
			sys.exit(-2)

		while True:
			unbrick(laptop)

	except SystemExit:
		pass
	except KeyboardInterrupt:
		print "Ctrl-C exit"
		pass
	except:
		print "Unexpected problem. Here's the details:"	
		traceback.print_exc()

	laptop.close()