#!/usr/bin/python3

#
# Simple tool to manipulate MicroDrive/Turbo CF Disk Images
# Bobbi 2021
#

import getopt
import os
import sys

# Sensible defaults for 512MB CF Card
default_cyls    = 1006
default_heads   = 16
default_sectors = 63
default_gsvers  = 0  # Apple //e

parttab = dict()
parttab['cyls']      = -1
parttab['heads']     = -1
parttab['sectors']   = -1
parttab['gsvers']    = -1
parttab['numparts1'] = 0
parttab['numparts2'] = 0
parttab['start1']    = list()
parttab['len1']      = list()
parttab['start2']    = list()
parttab['len2']      = list()

blocksz = 512

# This is 6502 code that must be present in the disk image from offset 0x0200 to 0x03ff
code_200 = [
  0x8e, 0x00, 0x82, 0x86, 0x43, 0x8a, 0x4a, 0x4a, 0x4a, 0x4a, 0x09, 0xc0, 0x85, 0x01, 0x8d, 0x81,
  0x80, 0xa0, 0x00, 0x84, 0x00, 0x84, 0x02, 0xa0, 0xff, 0xb1, 0x00, 0x8d, 0x80, 0x80, 0xa0, 0xf2,
  0xb1, 0x00, 0xc9, 0xc9, 0xd0, 0x38, 0xc8, 0xb1, 0x00, 0xc9, 0xc4, 0xd0, 0x31, 0xc8, 0xb1, 0x00,
  0xc9, 0xc5, 0xf0, 0x16, 0xc9, 0xc4, 0xf0, 0x0e, 0x85, 0x02, 0xc9, 0xd2, 0xf0, 0x0c, 0xc9, 0xd5,
  0xf0, 0x04, 0xc9, 0xd4, 0xd0, 0x18, 0xa9, 0x02, 0xd0, 0x02, 0xa9, 0x00, 0x8d, 0x02, 0x82, 0xa0,
  0xf9, 0xb1, 0x00, 0xc9, 0xca, 0xd0, 0x07, 0xc8, 0xb1, 0x00, 0xc9, 0xcc, 0xf0, 0x0a, 0x20, 0x58,
  0xfc, 0x20, 0x39, 0xfb, 0xa0, 0x08, 0xd0, 0x2e, 0x20, 0x5c, 0x81, 0xa0, 0x00, 0x84, 0x46, 0x84,
  0x47, 0x84, 0x44, 0xc8, 0x84, 0x42, 0xa9, 0x08, 0x85, 0x45, 0xad, 0x00, 0x82, 0x85, 0x43, 0x20,
  0xff, 0xff, 0xb0, 0x0d, 0xad, 0x00, 0x08, 0xc9, 0x01, 0xd0, 0x06, 0xae, 0x00, 0x82, 0x4c, 0x01,
  0x08, 0x20, 0x58, 0xfc, 0xa0, 0x06, 0x20, 0x36, 0x81, 0x4c, 0x00, 0xe0, 0xa8, 0x80, 0xc1, 0x80,
  0xd8, 0x80, 0xf1, 0x80, 0x08, 0x81, 0x1d, 0x81, 0xb0, 0x04, 0xad, 0xad, 0xa0, 0xd2, 0xc1, 0xcd,
  0xc6, 0xe1, 0xf3, 0xf4, 0xa0, 0xc9, 0xc4, 0xc5, 0xa0, 0xc3, 0xe1, 0xf2, 0xe4, 0xa0, 0xad, 0xad,
  0x00, 0xb1, 0x04, 0xad, 0xad, 0xa0, 0xd4, 0xf5, 0xf2, 0xe2, 0xef, 0xa0, 0xc9, 0xc4, 0xc5, 0xa0,
  0xc3, 0xe1, 0xf2, 0xe4, 0xa0, 0xad, 0xad, 0x00, 0xb0, 0x05, 0xa8, 0xe3, 0xa9, 0xa0, 0xb1, 0xb9,
  0xb9, 0xb6, 0xa0, 0xca, 0xef, 0xe1, 0xe3, 0xe8, 0xe9, 0xed, 0xa0, 0xcc, 0xe1, 0xee, 0xe7, 0xe5,
  0x00, 0xb1, 0x05, 0xd6, 0xef, 0xec, 0xf5, 0xed, 0xe5, 0xa0, 0xee, 0xef, 0xf4, 0xa0, 0xe6, 0xef,
  0xf2, 0xed, 0xe1, 0xf4, 0xf4, 0xe5, 0xe4, 0x00, 0xd0, 0x06, 0xc9, 0xc4, 0xc5, 0xa0, 0xc3, 0xe1,
  0xf2, 0xe4, 0xa0, 0xee, 0xef, 0xf4, 0xa0, 0xe6, 0xef, 0xf5, 0xee, 0xe4, 0x00, 0xd8, 0x07, 0xd6,
  0xef, 0xec, 0xf5, 0xed, 0xe5, 0xf3, 0xa0, 0xb1, 0xae, 0xae, 0xa0, 0xa0, 0xe1, 0xf6, 0xe1, 0xe9,
  0xec, 0xe1, 0xe2, 0xec, 0xe5, 0x00, 0xb9, 0x9c, 0x80, 0x85, 0x00, 0xb9, 0x9d, 0x80, 0x85, 0x01,
  0xa0, 0x00, 0xb1, 0x00, 0x8d, 0x55, 0x81, 0xc8, 0xb1, 0x00, 0x8d, 0x56, 0x81, 0xc8, 0xa2, 0x00,
  0xb1, 0x00, 0xf0, 0x07, 0x9d, 0xff, 0xff, 0xc8, 0xe8, 0xd0, 0xf5, 0x60, 0xac, 0x02, 0x82, 0x20,
  0x36, 0x81, 0xa0, 0x04, 0x20, 0x36, 0x81, 0xa5, 0x02, 0xf0, 0xf0, 0xac, 0x1a, 0xc8, 0xc0, 0x02,
  0x90, 0x11, 0xa5, 0x02, 0xf0, 0x0d, 0xb9, 0xec, 0x81, 0x48, 0xa0, 0x0a, 0x20, 0x36, 0x81, 0x68,
  0x8d, 0xe3, 0x07, 0xad, 0x61, 0xc0, 0x0d, 0x62, 0xc0, 0x10, 0xd0, 0xa5, 0x02, 0xf0, 0xf4, 0xad,
  0x00, 0xc0, 0x10, 0xef, 0xc9, 0xba, 0x90, 0x02, 0x29, 0xdf, 0xac, 0x10, 0xc0, 0xac, 0x1a, 0xc8,
  0xd9, 0xec, 0x81, 0xf0, 0x05, 0x88, 0xd0, 0xf8, 0xf0, 0xd9, 0x88, 0xf0, 0xae, 0x98, 0x0a, 0x0a,
  0xa8, 0x8c, 0x01, 0x82, 0xa2, 0x00, 0xb9, 0x80, 0xc8, 0x48, 0xbd, 0x80, 0xc8, 0x85, 0x1e, 0x68,
  0x9d, 0x80, 0xc8, 0xa5, 0x1e, 0x99, 0x80, 0xc8, 0xc8, 0xe8, 0xe0, 0x04, 0x90, 0xe8, 0xac, 0x01,
  0x82, 0xa2, 0x00, 0xb9, 0x40, 0xc8, 0x48, 0xbd, 0x40, 0xc8, 0x85, 0x1e, 0x68, 0x9d, 0x40, 0xc8,
  0xa5, 0x1e, 0x99, 0x40, 0xc8, 0xc8, 0xe8, 0xe0, 0x04, 0x90, 0xe8, 0x60, 0xb0, 0xb1, 0xb2, 0xb3,
  0xb4, 0xb5, 0xb6, 0xb7, 0xb8, 0xb9, 0xc1, 0xc2, 0xc3, 0xc4, 0xc5, 0xc6, 0xc7, 0x00, 0x00, 0x00
]


# Write val as a 16 bit integer into array data[], starting at index start
def write16(data, start, val):
	data[start] = val & 0x00ff
	data[start + 1] = (val & 0xff00) >> 8

# Read a 32 bit integer from array data[], starting at index start
def read32(data, start):
	return (data[start] + 256 * (data[start+ 1] 
	        + (256 * data[start+ 2] + (256 * data[start+ 3]))))

# Write val as a 32 bit integer into array data[], starting at index start
def write32(data, start, val):
	data[start] = val & 0x000000ff
	data[start + 1] = (val & 0x0000ff00) >> 8
	data[start + 2] = (val & 0x00ff0000) >> 16
	data[start + 3] = (val & 0xff000000) >> 24

# Parse partition table from file f and populate parttab
def read_part_tbl(f):
	with open(f, 'rb') as infile:
		b = infile.read(blocksz)
		parttab['cyls']    = b[0x02] + b[0x03] * 256
		parttab['heads']   = b[0x06] + b[0x07] * 256
		parttab['sectors'] = b[0x08] + b[0x09] * 256
		parttab['numparts1']  = b[0x0c]
		parttab['numparts2']  = b[0x0d]
		parttab['gsvers']  = b[0x18] + b[0x19] * 256
		parttab['start1'] = list()
		parttab['len1']   = list()
		parttab['start2'] = list()
		parttab['len2']   = list()
		if parttab['numparts1'] > 0:
			idx = 0x20
			for i in range(parttab['numparts1']):
				parttab['start1'].append(read32(b, idx + i * 4))
				parttab['len1'].append(read32(b, idx + 0x20 + i * 4))
		if parttab['numparts2'] > 0:
			idx = 0x80
			for i in range(parttab['numparts2']):
				parttab['start2'].append(read32(b, idx + i * 4))
				parttab['len2'].append(read32(b, idx + 0x20 + i * 4))

# Display partition table information in parttab
def print_part_tbl():
	if parttab['gsvers'] == 0:
		dma = 'iie'
	elif parttab['gsvers'] == 1:
		dma = 'gsrom01'
	elif parttab['gsvers'] == 2:
		dma = 'undefined'
	elif parttab['gsvers'] == 3:
		dma = 'gsrom03'
	elif parttab['gsvers'] == 4:
		dma = 'off'
	else:
		dma = '?'
	
	print('\nMicroDrive/Turbo Partition Table\n')
	print('  Cylinders:   {}'.format(parttab['cyls']))
	print('  Heads:       {}'.format(parttab['heads']))
	print('  Sectors:     {}'.format(parttab['sectors']))
	print('  DMA Mode:    {}'.format(dma))
	if parttab['numparts1'] > 0:
		print('\n  #    Start      End   Length')
		for i in range(parttab['numparts1']):
			st = parttab['start1'][i]
			ln = parttab['len1'][i]
			print('  {0} {1:8d} {2:8d} {3:8d}'
			      .format(i + 1, st, st + ln - 1, ln))
	if parttab['numparts2'] > 0:
		for i in range(parttab['numparts2']):
			st = parttab['start2'][i]
			ln = parttab['len2'][i]
			print('  {0} {1:8d} {2:8d} {3:8d}'
			      .format(i + 9, st, st + ln - 1, ln))

# Write partition table in parttab to file f
# Writes 256 blocks (ie 128K)
def write_part_tbl(f):
	magic = 0xccca
	if ((parttab['cyls'] == -1)
	    or (parttab['heads'] == -1)
	    or (parttab['sectors'] == -1)
	    or (parttab['gsvers'] == -1)
	    or (parttab['numparts1'] == -1)
	    or (parttab['numparts2'] == -1)):
		print('** Not writing partition table - uninitialized data **')
		exit()
	b = [0] * blocksz
	write16(b, 0x00, magic)
	write16(b, 0x02, parttab['cyls'])
	write16(b, 0x06, parttab['heads'])
	write16(b, 0x08, parttab['sectors'])
	b[0x0c] = parttab['numparts1']
	b[0x0d] = parttab['numparts2']
	write16(b, 0x18, parttab['gsvers'])
	for i in range(parttab['numparts1']):
		write32(b, 0x20 + i * 4, parttab['start1'][i])
		write32(b, 0x40 + i * 4, parttab['len1'][i])
	for i in range(parttab['numparts2']):
		write32(b, 0x80 + i * 4, parttab['start2'][i])
		write32(b, 0xa0 + i * 4, parttab['len2'][i])
	with open(f, 'wb') as outfile:
		block = bytes(b)
		outfile.write(block)
		block = bytes(code_200)
		outfile.write(block)
		block = bytes([0] * blocksz)
		for i in range(254):
			outfile.write(block)

# Obtain volume name of a partition
def get_volname(infile, startidx):
	infile.seek(startidx * blocksz + 1024)
	mdb = infile.read(2)
	if (mdb[0] == 0x42) and (mdb[1] == 0x44): # HFS Signature in Block 2
		fstype = 'HFS'
		infile.seek(startidx * blocksz + 1024 + 36)
		volbytes = infile.read(27)
		volname = bytearray()
		for i in range(0, volbytes[0] & 0x0f):
			volname.append(volbytes[i + 1])
	else:
		fstype = 'ProDOS' # If not HFS, must be ProDOS
		infile.seek(startidx * blocksz + 1024 + 4)
		volbytes = infile.read(15)
		volname = bytearray()
		for i in range(0, volbytes[0] & 0x0f):
			volname.append(volbytes[i + 1])
	return (fstype, volname.decode('ascii'))

# Extract partition n from file f and write it to file named
# according to the volume name
def extract_partition(f, n):
	if n < 8:
		st = parttab['start1'][n]
		ln = parttab['len1'][n]
	else:
		st = parttab['start2'][n - 8]
		ln = parttab['len2'][n - 8]
	with open(f, 'rb') as infile:
		(fstype, volname) = get_volname(infile, st)
		print('** Partition ' + str(n + 1) + ' ' + fstype + ' Volume \'' + volname + '\'')
		if fstype == 'HFS':
			outputfile = volname + '.hfs'
		elif fstype == 'ProDOS':
			outputfile = volname + '.po'
		infile.seek(st * blocksz)
		d = infile.read(ln * blocksz)
		print('Writing {} bytes to {}'.format(len(d), outputfile))
		with open(outputfile, 'wb') as outfile:
			outfile.write(d)

# Extract all partitions from file f
def explode_mdt_image(f):
	for i in range(parttab['numparts1']):
		extract_partition(f, i)
	for i in range(parttab['numparts2']):
		extract_partition(f, i + 8)

# Build an MDT image file f from a list of .po files and update the
# partition table to match the files
def assemble_mdt_image(f, files):

	# Add a partition to parttab. i is the zero-based index
	def add_partition(i, start, length):
		if i < 8:
			parttab['start1'].append(start)
			parttab['len1'].append(length)
			parttab['numparts1'] += 1
		else:
			parttab['start2'].append(start)
			parttab['len2'].append(length)
			parttab['numparts2'] += 1

	currblk = 256
	numparts = 0
	for filename in files:
		status = os.stat(filename)
		sz = status.st_size // 512
		add_partition(numparts, currblk, sz)
		numparts += 1
		currblk += sz
	write_part_tbl(f)
	for filename in files:
		with open(filename, 'rb') as infile:
			print('  Reading {} ...'.format(filename))
			d = infile.read()
			with open(f, 'ab') as outfile:
				outfile.write(d)
	print('New CF card image created in {}'.format(f))

def initialize_blank_mdt():
	pass

def usage():
	print('Usage:')
	print('  mdttool [-haxl] [-C nnnn] [-H nn] [-S nn] [-D mode] [-o outputfile] inputfile [inputfile2...]\n')
	print('  -h       Show this help')
	print('  -a       Assemble CF Card image from list of .PO disk images')
	print('  -x       eXplode a CF Card image into constituent partitions')
	print('  -l       Display partition table');
	print('  -C  nnnn Set number of cylinders (when assembling new image)')
	print('  -H  nn   Set number of heads (when assembling new image)')
	print('  -S  nn   Set number of sectors (when assembling new image)')
	print('  -D  mode DMA mode (when assembling new image)')
	print('           mode can be off, iie, gsrom01, gsrom03')

# Main entry point and argument handling
def main():
	try:
		opts, args = getopt.getopt(
		               sys.argv[1:],
		               'ho:axlC:H:S:D:',
		               ['help', 'output=', 'assemble', 'explode', 'list',
		                'cylinders', 'heads', 'sectors', 'dma'])
	except getopt.GetoptError as err:
		print(err)
		usage()
		sys.exit(2)

	infile = None
	outfile = None
	assemble = False
	explode = False
	printtab = False
	cylinders = default_cyls
	heads = default_heads
	sectors = default_sectors
	gsvers = default_gsvers

	for o, a in opts:
		if o in ('-h', '--help'):
			usage()
			sys.exit()
		elif o in ('-o', '--output'):
			outfile = a
		elif o in ('-a', '--assemble'):
			assemble = True
		elif o in ('-x', '--explode'):
			explode = True
		elif o in ('-l', '--list'):
			printtab = True
		elif o in ('-C', '--cylinders'):
			cylinders = int(a)
		elif o in ('-H', '--heads'):
			heads = int(a)
		elif o in ('-S', '--sectors'):
			sectors = int(a)
		elif o in ('-D', '--dma'):
			if a == 'iie':
				gsvers = 0
			elif a == 'gsrom01':
				gsvers = 1
			elif a == 'gsrom03':
				gsvers = 3
			elif a == 'off':
				gsvers = 4
			else:
				print('Bad mode {}'.format(a))
				usage()
				sys.exit(2)
		else:
			assert False, 'Unhandled option'

	if assemble:
		if len(args) == 0:
			print('Need at least one input file to assemble')
			usage()
			sys.exit(2)
		if outfile == None:
			print('Must specify output file with -o')
			usage()
			sys.exit(2)
		parttab['cyls'] = cylinders
		parttab['heads'] = heads
		parttab['sectors'] = sectors
		parttab['gsvers'] = gsvers
		assemble_mdt_image(outfile, args)
	else:
		if len(args) == 0:
			print('No input file provided')
			usage()
			sys.exit(2)
		if len(args) > 1:
			print('Provide only one input file')
			usage()
			sys.exit(2)
		read_part_tbl(args[0])
		if printtab:
			print_part_tbl()
		if explode:
			explode_mdt_image(args[0])

	sys.exit()
	
		
main()