ftdi-cbus-ctrlep-tdx.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import time
import usb

ftdi_eval = {
    'tty': "/dev/ttyUSB3",
    'serial': "A105N7O2",
};

BITMODE_CBUS = 0x20

SIO_SET_BITMODE_REQUEST = 0x0b

# FTDIs CBUS bitmode expect the following value:
# CBUS Bits
# 3210 3210
#      |------ Output Control 0->LO, 1->HI
# |----------- Input/Output   0->Input, 1->Output

#
# This script assumes:
# - CBUS3 connected to RESET_EXT# 
# - CBUS2 connected to OE# (recovery mode)
#

# PyUSB control endpoint communication, see also:
# https://github.com/pyusb/pyusb/blob/master/docs/tutorial.rst
#
# To use the script from a screen using Ctrl+A, [letter from below], use the following
# bindings in ~/.screenrc
# bind n exec ctrlep.py b
# bind r exec ctrlep.py r
# bind i exec ctrlep.py i


def ftdi_set_bitmode(dev, bitmask):
    bmRequestType = usb.util.build_request_type(usb.util.CTRL_OUT,
                                                usb.util.CTRL_TYPE_VENDOR,
                                                usb.util.CTRL_RECIPIENT_DEVICE)

    wValue = bitmask | (BITMODE_CBUS << BITMODE_CBUS)
    dev.ctrl_transfer(bmRequestType, SIO_SET_BITMODE_REQUEST, wValue)


def module_on(dev):
    # Set CBUS3 tristate, module run...
    ftdi_set_bitmode(dev, 0x00)

def module_off(dev):
    # Set CBUS3 low, module in reset...
    ftdi_set_bitmode(dev, 0x80)

def module_reboot(dev):
    module_off(dev)
    time.sleep(0.1)
    module_on(dev)

def module_recovery(dev, recvoery_high_active):
    # Set CBUS2/3 low, and then CBUS3 tristate
    recovery_bit = 0x4 if recvoery_high_active else 0x0
    ftdi_set_bitmode(dev, 0xC0 | recovery_bit)
    time.sleep(0.1)
    ftdi_set_bitmode(dev, 0x40 | recovery_bit)
    time.sleep(0.2)
    ftdi_set_bitmode(dev, 0x00)

def do_cmd(dev, cmd):
    if cmd == '0':
        module_off(dev)
    elif cmd == '1':
        module_on(dev)
    elif cmd == 'b':
        module_reboot(dev)
    elif cmd == 'r':
        module_recovery(dev, False)
    elif cmd == 'i':
        module_recovery(dev, True)
    else:
        return False
    return True

def get_tty():
    # Find FTDI using Environment variable set by screen
    sty = os.environ['STY']
    pid = sty.split('.')[0]
    with open("/proc/{}/cmdline".format(pid), 'r') as f:
        for line in f.readline().split('\0'):
            if line.startswith("/dev/"):
                return line

def get_serial_from_tty(ttyname):
    p = os.popen("udevadm info {} -q property".format(ttyname), "r")
    while True:
        line = p.readline()
        if not line:
            raise Exception("Serial number of tty {} not found".format(ttyname))
        if line.startswith("ID_SERIAL_SHORT"):
            return line.split('=')[1]

def main():
    """Main program"""
    restart = True

    if len(sys.argv) > 1 and sys.argv[1].startswith("/dev/"):
        ttyname = sys.argv[1]
    else:
        # Get tty name when calling from screen...
        ttyname = get_tty()

    # Find FTDI serial number of that tty...
    serial = get_serial_from_tty(ttyname).strip()

    dev = usb.core.find(custom_match = \
            lambda d: \
                d.idVendor==0x0403 and
                d.idProduct==0x6001 and
                d.serial_number==serial)

    if len(sys.argv) > 1:
        if len(sys.argv[-1]) == 1:
            if do_cmd(dev, sys.argv[-1]):
                return
        print("Only [r/b/0/1] are valid commands")
        return

    while restart:
        try:
            newcmd = input('Command (recovery/boot/off/on)? [r/b/0/1] ')
            if newcmd is not None:
                do_cmd(dev, newcmd)
        except KeyboardInterrupt:
            module_on(dev)
            break

main()