#!/usr/bin/python3.1

# scans for BusPirate port, and reads 4Mbit flash image
# NOTE: Will overwrite target file without warning!

# Version 1.0  11.12.2012

# Copyright (C) 2012, Arno Wagner <arno@wagner.name>
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# version 2, or a later version at your choice, as published by the
# Free Software Foundation.
#
# 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, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301


import serial, time, sys

dumpfile='flash_data.img'

candidates = ['/dev/ttyUSB0','/dev/ttyUSB1','/dev/ttyUSB2','/dev/ttyUSB3',
        '/dev/ttyUSB4','/dev/ttyUSB5','/dev/ttyUSB6','/dev/ttyUSB7',
        '/dev/ttyUSB8','/dev/ttyUSB9','/dev/ttyUSB10','/dev/ttyUSB11',
        '/dev/ttyUSB12','/dev/ttyUSB13','/dev/ttyUSB14','/dev/ttyUSB15']



def test_if(name):
  global wait_time
  try:
    s = serial.Serial(port=name, baudrate=115200, timeout=0.01)
  except Exception as e:
    return None
  else:
    return s

    
def test_bp(s):
  if s == None:
    return None 
  else:  
    s.write(b'#\n')
    i = s.read(10000)
    sl = i.splitlines() 

    if len(sl) < 4 or \
       ( sl[2][0:10] != b'Bus Pirate' and \
         sl[3][0:10] != b'Bus Pirate' ):
      return None;
    return True


def find_bp():
  # tests list of candidates. Returns serial instance if exactly one found. 
  # If several found, aborts with exception.
  # If none found, returns None           
  cnt = 0
  s_bp = None
  n_bp = None
  n_str = ''
  for n in candidates:
    s = test_if(n)
    is_bp = test_bp(s)
    if is_bp: 
      cnt += 1
      s_bp = s
      n_bp = n
      n_str += n + '\n'
  if cnt > 1:
    raise IOError('More than one Bus Pirate found! Interfaces: \n'+n_str)
  return (s_bp, n_bp)  
    
    
def adr24(n):
  # takes an int and returns a 24 but hex address MSB,...,LSB
  # e.g. 0x01 0x00 0x23 
  b1 = n//65535
  b2 = (n//256)%256
  b3 = n%256
  s = '0x{:02X} '.format(b1) + '0x{:02X} '.format(b2) + '0x{:02X}'.format(b3)     
  print('adr24: ',n,' -> ', s)
  return(s.encode('ascii'))  
    
    
def safe_exit(msg):
  # exit with BP reset
  print('  ** ERROR -- safe_exit() called **')
  print('  MSG: ',msg)
  s.write(b'#\n')
  i = s.read(100)
#  sl = i.decode('ascii').splitlines()
#  print(sl)
  # Note: seems reset from SPI does sometimes not echo the #...
#  if sl[0]!= 'RESET' and sl[1] != 'RESET': 
#    print('  ** error exit reset failed! **')
  sys.exit(-1)  
  

a = find_bp()
(s,n) = a
#print('device:    ', n)
#print('seri_inst: ', s) 
 
# Start interaction
s = serial.Serial(port=n, baudrate=115200, timeout=0.01)

# reset just in case
s.write(b'#\n')
i = s.read(200)
#sl = i.decode('ascii').splitlines()
#print(sl)
#if sl[1] != 'RESET': safe_exit('initiial reset failed')

# m set SPI
s.write(b'm\n') # 'm'
i = s.read(100)
sl = i.decode('ascii').splitlines()
#print(sl)
    
# SPI
s.write(b'5\n') # Menu item 5 = SPI
i = s.read(100)
sl = i.decode('ascii').splitlines()
#print(sl)

s.write(b'4\n') # 1=30khz, 2=125kHz, 3=250kHz, 4=1MHz
i = s.read(100)
sl = i.decode('ascii').splitlines()
#print(sl)

s.write(b'1\n') # 1= clock idle low (default)  2: clock idle high
i = s.read(100)
sl = i.decode('ascii').splitlines()
#print(sl)

s.write(b'2\n') # 1:idle to active 2:active to idle (default)
# note: for the 25x40 the datasheet seems to indicate 1 is
# correct, but 2 works also. But for a a25l040, 1 fails while
# 2 works. So likely 2 is really the usually correct choice.
i = s.read(100)
sl = i.decode('ascii').splitlines()
#print(sl)

s.write(b'1\n') # sample signal 1: middle  2:end
i = s.read(100)
sl = i.decode('ascii').splitlines()
#print(sl)

s.write(b'2\n') # 1:CS 2:/CS
i = s.read(100)
sl = i.decode('ascii').splitlines()
#print(sl)

s.write(b'2\n') # 1: H=Hi-Z 2: H=3v3 Note: Pullups are problematic
i = s.read(100)
sl = i.decode('ascii').splitlines()
#print(sl)
if sl[1] != 'Ready': safe_exit('SPI detail setting failed')

s.write(b'W\n') # Activate power
i = s.read(100)
sl = i.decode('ascii').splitlines()
#print(sl)
if sl[1] != 'Power supplies ON' and \
   sl[1] != 'POWER SUPPLIES ON':
     safe_exit('power ON failed')

# give is a bit of time to come up and settle
time.sleep(0.200)

# read data
outf = open(dumpfile, 'wb')

read_size = 128

for adr in range(0, 0x80000, read_size):
  address = adr24(adr)
  cmd = b'[0x03 ' + address + b' r:'+str(read_size).encode('ascii')+b']\n'
  print('cmd: ', cmd)
  s.write(cmd) # 
  i = s.read(2000)
  sl = i.decode('ascii').splitlines()
  #print(sl)
  r = sl[6]
  if r[0:6] != 'READ: ': safe_exit('read failed')

  ds = r[6:]
#  print('ds: ',ds)
  dss = ds.split()
#  for n in dss: print(int(n, 16), end=' ')
#  print()
  print('len: ', len(dss))
  for n in dss: 
    c = int(n, 16)
    b = bytes([c]) 
    outf.write(b) 
outf.close()


# Final reset
s.write(b'#\n')
i = s.read(200)
#sl = i.decode('ascii').splitlines()
#print(sl)
# Note: Seems reset from SPI does sometimes not echo the '#'
#if sl[0] != 'RESET' and sl[1] != 'RESET': safe_exit('final reset failed')