#!/usr/bin/env python2
# -*- coding: utf-8 -*-
# Copyright 2010, 2011 Michael Ossmann
# Copyright 2015 Travis Goodspeed
#
# This file was forked from Project Ubertooth as a DFU client for the
# TYT MD380, an amateur radio for the DMR protocol on the UHF bands.
# This script implements a lot of poorly understood extensions unique
# to the MD380.

import sys
import usb.core
import struct
import time

# The tricky thing is that *THREE* different applications all show up
# as this same VID/PID pair.
#
# 1. The Tytera application image.
# 2. The Tytera bootloader at 0x08000000
# 3. The mask-rom bootloader from the STM32F405.

class Enumeration(object):
    def __init__(self, id, name):
        self._id = id
        self._name = name
        setattr(self.__class__, name, self)
        self.map[id] = self

    def __int__(self):
        return self.id

    def __repr__(self):
        return self.name

    @property
    def id(self):
        return self._id

    @property
    def name(self):
        return self._name

    @classmethod
    def create_from_map(cls):
        for id, name in cls.map.items():
            cls(id, name)

class Request(Enumeration):
    map = {
        0: 'DETACH',
        1: 'DNLOAD',
        2: 'UPLOAD',
        3: 'GETSTATUS',
        4: 'CLRSTATUS',
        5: 'GETSTATE',
        6: 'ABORT',
    }
Request.create_from_map()

class State(Enumeration):
    map = {
        0: 'appIDLE',
        1: 'appDETACH',
        2: 'dfuIDLE',
        3: 'dfuDNLOAD_SYNC',
        4: 'dfuDNBUSY',
        5: 'dfuDNLOAD_IDLE',
        6: 'dfuMANIFEST_SYNC',
        7: 'dfuMANIFEST',
        8: 'dfuMANIFEST_WAIT_RESET',
        9: 'dfuUPLOAD_IDLE',
        10: 'dfuERROR',
    }
State.create_from_map()

class Status(Enumeration):
    map = {
        0x00: 'OK',
        0x01: 'errTARGET',
        0x02: 'errFILE',
        0x03: 'errWRITE',
        0x04: 'errERASE',
        0x05: 'errCHECK_ERASED',
        0x06: 'errPROG',
        0x07: 'errVERIFY',
        0x08: 'errADDRESS',
        0x09: 'errNOTDONE',
        0x0A: 'errFIRMWARE',
        0x0B: 'errVENDOR',
        0x0C: 'errUSBR',
        0x0D: 'errPOR',
        0x0E: 'errUNKNOWN',
        0x0F: 'errSTALLEDPKT',
    }
Status.create_from_map()

verbose = False

# Detach from the DFU target.
def dfu_detach(dev):
    dev.ctrl_transfer(0x21, Request.DETACH, 0, 0, None)

# Convert a byte from BCD to integer.
def bcd(b):
    return int("%02x" % b)

def dfu_download(dev, block_number, data):
    dev.ctrl_transfer(0x21, Request.DNLOAD, block_number, 0, data)
    # time.sleep(0.1);

def dfu_set_address(dev, address):
    a = address & 0xFF
    b = (address >> 8) & 0xFF
    c = (address >> 16) & 0xFF
    d = (address >> 24) & 0xFF
    dev.ctrl_transfer(0x21, Request.DNLOAD, 0, 0, [0x21, a, b, c, d])
    dfu_get_status(dev)  # this changes state
    status = dfu_get_status(dev)  # this gets the status
    if status[2] == State.dfuDNLOAD_IDLE:
        if verbose:
            print("Set pointer to 0x%08x." % address)
        dfu_enter_dfu_mode(dev)
    else:
        if verbose:
            print("Failed to set pointer.")
        return False
    return True

def dfu_erase_block(dev, address):
    a = address & 0xFF
    b = (address >> 8) & 0xFF
    c = (address >> 16) & 0xFF
    d = (address >> 24) & 0xFF
    dev.ctrl_transfer(0x21, Request.DNLOAD, 0, 0, [0x41, a, b, c, d])
    # time.sleep(0.5);
    dfu_get_status(dev)  # this changes state
    status = dfu_get_status(dev)  # this gets the status
    if status[2] == State.dfuDNLOAD_IDLE:
        if verbose:
            print("Erased 0x%08x." % address)
        dfu_enter_dfu_mode(dev)
    else:
        if verbose:
            print("Failed to erase block.")
        return False
    return True

# Sends a secret MD380 command.
def dfu_md380_custom(dev, a, b):
    a &= 0xFF
    b &= 0xFF
    dev.ctrl_transfer(0x21, Request.DNLOAD, 0, 0, [a, b])
    dfu_get_status(dev)  # this changes state
    time.sleep(0.1)
    status = dfu_get_status(dev)  # this gets the status
    if status[2] == State.dfuDNLOAD_IDLE:
        if verbose:
            print("Sent custom %02x %02x." % (a, b))
        dfu_enter_dfu_mode(dev)
    else:
        print("Failed to send custom %02x %02x." % (a, b))
        return False
    return True

def dfu_md380_reboot(dev):
    """Sends the MD380's secret reboot command."""
    a = 0x91
    b = 0x05
    dev.ctrl_transfer(0x21, Request.DNLOAD, 0, 0, [a, b])
    try:
        dfu_get_status(dev)  # this changes state
    except:
        pass
    return True

def dfu_upload(dev, block_number, length, index=0):
    if verbose:
        print("Fetching block 0x%x." % block_number)
    data = dev.ctrl_transfer(0xA1,              # request type
                             Request.UPLOAD,    # request
                             block_number,      # wValue
                             index,             # index
                             length)            # length
    return data

def dfu_get_command(dev):
    data = dev.ctrl_transfer(0xA1,              # request type
                             Request.UPLOAD,    # request
                             0,                 # wValue
                             0,                 # index
                             32)                # length
    dfu_get_status(dev)
    return data

def dfu_get_status(dev):
    status_packed = dev.ctrl_transfer(0xA1, Request.GETSTATUS, 0, 0, 6)
    status = struct.unpack('<BBBBBB', status_packed)
    return (Status.map[status[0]], (((status[1] << 8) | status[2]) << 8) | status[3],
            State.map[status[4]], status[5])

def dfu_clear_status(dev):
    dev.ctrl_transfer(0x21, Request.CLRSTATUS, 0, 0, None)

def dfu_get_state(dev):
    state_packed = dev.ctrl_transfer(0xA1, Request.GETSTATE, 0, 0, 1)
    return State.map[struct.unpack('<B', state_packed)[0]]

def dfu_abort(dev):
    dev.ctrl_transfer(0x21, Request.ABORT, 0, 0, None)

def dfu_wait_till_ready(dev, desired_state=State.dfuIDLE):
    state = 11
    status, timeout, state, discarded = dfu_get_status(dev)
    while state != State.dfuIDLE:
        dfu_clear_status(dev)
        status, timeout, state, discarded = dfu_get_status(dev)
    return

def dfu_enter_dfu_mode(dev):
    action_map = {
        State.dfuDNLOAD_SYNC: dfu_abort,
        State.dfuDNLOAD_IDLE: dfu_abort,
        State.dfuMANIFEST_SYNC: dfu_abort,
        State.dfuUPLOAD_IDLE: dfu_abort,
        State.dfuERROR: dfu_clear_status,
        State.appIDLE: dfu_detach,
        State.appDETACH: dfu_wait,
        State.dfuDNBUSY: dfu_wait,
        State.dfuMANIFEST: dfu_abort,
        State.dfuMANIFEST_WAIT_RESET: dfu_wait,
        State.dfuIDLE: dfu_wait
    }

    try:
        while True:
            state = dfu_get_state(dev)
            if state == State.dfuIDLE:
                break
            action = action_map[state]
            action(dev)
    except usb.core.USBError as e:
        print("Failed to enter DFU mode.")
        if len(e.args) > 0 and e.args[0] == 'Pipe error':
            print("Is bootloader running?")
        exit(1)

def dfu_wait(dev):
    time.sleep(0.1)

def dfu_widestr(str):
    tr = ""
    for c in str:
        tr = tr + c + "\0"
    return tr + "\0\0"

def hexdump(string):
    """God awful hex dump function for testing."""
    buf = ""
    i = 0
    for c in string:
        buf += "%02x" % c
        i += 1
        if i & 3 == 0:
            buf += " "
        if i & 0xf == 0:
            buf += "   "
        if i & 0x1f == 0:
            buf += "\n"

    print(buf)

#
# Parse command line.
#
if len(sys.argv) != 2:
    print("Usage: md380-read <codeplug.bin>")
    exit(1)
filename = sys.argv[1]

#
# Open the device.
#
dev = usb.core.find(idVendor = 0x0483, idProduct = 0xdf11)
if dev is None:
    raise RuntimeError('Device not found')
dev.set_interface_altsetting(interface=0, alternate_setting=0)
dev.default_timeout = 3000

#
# Enable DFU mode.
#
dfu_enter_dfu_mode(dev)

#
# Enter Programming Mode.
#
dfu_md380_custom(dev, 0x91, 0x01)

#
# Upload a codeplug from the radio to the host.
#
dfu_set_address(dev, 0x00000000)  # Zero address, used by configuration tool.

f = open(filename, 'wb')
block_size = 1024
try:
    # Codeplug region is 0 to 3ffffff, but only the first 256k are used.
#    for block_number in range(2, 2+256*8):
    for block_number in range(2, 256*4*16+2):
        data = dfu_upload(dev, block_number, block_size)
        status, timeout, state, discarded = dfu_get_status(dev)

        # print("Status is: %x %x %x %x" % (status, timeout, state, discarded))
        sys.stdout.write('.')
        sys.stdout.flush()
        if len(data) == block_size:
            f.write(data)
            # hexdump(data);
        else:
            raise Exception('Upload failed to read full block.  Got %i bytes.' % len(data))
            # dfu_md380_reboot(dev)
finally:
    print("Done.")

print('Read complete')