xdftool

#!/usr/bin/env python2.7
# xdftool
# swiss army knife for adf and hdf amiga disk images

import sys
import argparse
import os.path

from amitools.fs.ADFSVolume import ADFSVolume
from amitools.fs.blkdev.BlkDevFactory import BlkDevFactory
from amitools.fs.FSError import *
from amitools.fs.Imager import Imager
from amitools.fs.Repacker import Repacker
from amitools.fs.block.BootBlock import BootBlock
from amitools.fs.block.RootBlock import RootBlock
from amitools.util.CommandQueue import CommandQueue
from amitools.util.HexDump import *
from amitools.util.DisAsm import DisAsm
import amitools.util.KeyValue as KeyValue
from amitools.fs.rdb.RDisk import RDisk
from amitools.fs.FSString import FSString

# system encoding
def make_fsstr(s):
  return FSString(s.decode(sys.stdout.encoding))

# ----- commands -----
class Command:
  def __init__(self, args, opts, edit=False):
    self.args = args
    self.opts = opts
    self.edit = edit
    self.exit_code = 0
    
    self.volume = None
    self.blkdev = None
    
  def run(self, blkdev, vol):
    # optional init blkdev function
    if hasattr(self, "init_blkdev"):
      if blkdev == None:
        self.blkdev = self.init_blkdev(self.args.image_file)
        if self.blkdev == None:
          return 5
        blkdev = self.blkdev
        
    # optional init volume function
    if hasattr(self, "init_vol"):
      # close old
      if vol != None:
        vol.close()
      # create new volume
      self.volume = self.init_vol(blkdev)
      if self.volume == None:
        return 6
      vol = self.volume

    # common handler
    if hasattr(self, 'handle_blkdev'):
      return self.handle_blkdev(blkdev)
    elif hasattr(self, 'handle_vol'):
      return self.handle_vol(vol)
    else:
      return 0
  
  def has_init_blkdev(self):
    return hasattr(self, 'init_blkdev')
    
  def need_volume(self):
    return hasattr(self, 'handle_vol') and not hasattr(self, 'init_vol')

# ----- command handler -----
class FSCommandQueue(CommandQueue):
  def __init__(self, args, cmd_list, sep, cmd_map):
    CommandQueue.__init__(self, cmd_list, sep, cmd_map)
    self.args = args
    self.blkdev = None
    self.volume = None
  
  def run(self):
    self.img = self.args.image_file
    try:
      # main command loop
      exit_code = CommandQueue.run(self)
    except FSError as e:
      cmd = "'%s'" % " ".join(self.cmd_line)
      print cmd,"FSError:",e
      exit_code = 3
    except IOError as e:
      cmd = "'%s'" % " ".join(self.cmd_line)
      print cmd,"IOError:",e
      exit_code = 4
    finally:
      # close volume
      if self.volume != None:
        self.volume.close()
        if self.args.verbose:
          print "closing volume:",self.img
      # close blkdev
      if self.blkdev != None:
        self.blkdev.close()
        if self.args.verbose:
          print "closing image:",self.img
    return exit_code

  def create_cmd(self, cclass, name, opts):
    return cclass(self.args, opts)

  def _open_volume(self):
    # setup volume
    if self.volume == None:
      self.volume = ADFSVolume(self.blkdev)
      if self.args.verbose:
        print "opening volume:",self.img
      self.volume.open()
  
  def run_first(self, cmd_line, cmd):
    self.cmd_line = cmd_line
    
    # check if first command is an init command
    if not cmd.has_init_blkdev():
      # auto add 'open' command
      pre_cmd = OpenCmd(self.args, [])
      if self.args.verbose:
        print "auto open command:",self.cmd_line
      exit_code = pre_cmd.run(self.blkdev, self.volume)
      if self.args.verbose:
        print "auto open exit_code:",exit_code
      if exit_code != 0:
        return exit_code
      self.blkdev = pre_cmd.blkdev
      # setup volume (if necessary)
      if cmd.need_volume():
        self._open_volume()
        
    # run first command    
    if self.args.verbose:
      print "command:",self.cmd_line
    if cmd.edit and self.args.read_only:
      raise IOError("Edit commands not allowed in read-only mode")

    # check code of command after __init__ parsing
    if cmd.exit_code != 0:
      return cmd.exit_code  
    
    # perform command
    exit_code = cmd.run(self.blkdev, self.volume)
    if cmd.blkdev != None:
      self.blkdev = cmd.blkdev
    if cmd.volume != None:
      self.volume = cmd.volume

    # final exit code
    if self.args.verbose:
      print "exit_code:",exit_code
    return exit_code
    
  def run_next(self, cmd_line, cmd):
    self.cmd_line = cmd_line
    if self.args.verbose:
      print "command:",self.cmd_line
    # verify command
    if cmd.edit and self.args.read_only:
      raise IOError("Edit commands not allowed in read-only mode")
    # make sure volume is set up
    if self.volume == None and cmd.need_volume():
      self._open_volume()
    # run command
    exit_code = cmd.run(self.blkdev, self.volume)
    if cmd.blkdev != None:
      self.blkdev = cmd.blkdev
    if cmd.volume != None:
      self.volume = cmd.volume
    if self.args.verbose:
      print "exit_code:",exit_code
    return exit_code

# ----- Init: Open/Create -----

class OpenCmd(Command):
  def init_blkdev(self, image_file):
    opts = KeyValue.parse_key_value_strings(self.opts)
    f = BlkDevFactory()
    return f.open(image_file, options=opts, read_only=args.read_only)
    
class CreateCmd(Command):
  def __init__(self, args, opts):
    Command.__init__(self, args, opts, edit=True)
    
  def init_blkdev(self, image_file):
    opts = KeyValue.parse_key_value_strings(self.opts)
    f = BlkDevFactory()
    return f.create(image_file, options=opts, force=args.force)

class FormatCmd(Command):
  def init_blkdev(self, image_file):
    opts = KeyValue.parse_key_value_strings(self.opts[1:])
    f = BlkDevFactory()
    return f.create(image_file, options=opts, force=args.force)

  def init_vol(self, blkdev):
    vol = ADFSVolume(blkdev)
    n = len(self.opts)
    if n == 0:
      print "Usage: format <volume_name> [ffs] [intl] [dircache]"
      return None
    else:
      is_ffs = "ffs" in self.opts
      is_intl = "intl" in self.opts
      is_dircache = "dircache" in self.opts
      vol_name = make_fsstr(self.opts[0])
      vol.create(vol_name, is_ffs=is_ffs, is_intl=is_intl, is_dircache=is_dircache)
      return vol

# ----- Pack/Unpack -----

class PackCmd(Command):
  def __init__(self, args, opts):
    Command.__init__(self, args, opts, edit=True)
    self.imager = Imager()
    n = len(self.opts)
    if n == 0:
      print "Usage: pack <in_path> [out_size]"
      self.exit_code = 1
    else:
      self.in_path = self.opts[0]
      if n > 1:
        self.blkdev_opts = KeyValue.parse_key_value_strings(opts[1:])
      else:
        self.blkdev_opts = None
      self.imager.pack_begin(self.in_path)
      
  def init_blkdev(self, image_file):
    return self.imager.pack_create_blkdev(self.in_path, image_file, force=self.args.force, options=self.blkdev_opts)
  
  def init_vol(self, blkdev):
    return self.imager.pack_create_volume(self.in_path, blkdev)
  
  def handle_vol(self, volume):
    self.imager.pack_root(self.in_path, volume)
    self.imager.pack_end(self.in_path, volume)
    if self.args.verbose:
      print "Packed %d bytes" % (self.imager.get_total_bytes())
    return 0

class UnpackCmd(Command):
  def handle_vol(self, vol):
    n = len(self.opts)
    if n == 0:
      print "Usage: unpack <out_path>"
      return 1
    else:
      out_path = self.opts[0]
      img = Imager()
      img.unpack(vol, out_path)
      if self.args.verbose:
        print "Unpacked %d bytes" % (img.get_total_bytes())
      return 0

class RepackCmd(Command):
  def __init__(self, args, opts):
    Command.__init__(self, args, opts, edit=True)
    n = len(self.opts)
    if n == 0:
      print "Usage: repack <src_path> [in_size]"
      self.exit_code = 1
    in_img = self.opts[0]
    in_opts = KeyValue.parse_key_value_strings(self.opts[1:])
    self.repacker = Repacker(in_img, in_opts)
    if not self.repacker.create_in():
      self.exit_code = 2

  def init_blkdev(self, image_file):
    return self.repacker.create_out_blkdev(image_file)
    
  def init_vol(self, blkdev):
    return self.repacker.create_out_volume(blkdev)
  
  def handle_vol(self, vol):
    self.repacker.repack()
    return 0

# ----- Query Image -----

# list: list directory tree
class ListCmd(Command):
  def handle_vol(self, vol):
    n = len(self.opts)
    if n == 0:
      vol.root_dir.list(all=True)
      show_info = True
      show_all = True
      node = vol.get_root_dir()
    else:
      name = make_fsstr(self.opts[0])
      node = vol.get_path_name(name)
      if node != None:
        show_all = "all" in self.opts
        show_detail = "detail" in self.opts
        show_info = "info" in self.opts
        node.list(all=show_all, detail=show_detail)
      else:
        print "ERROR path not found:",node
        return 2
    if show_info:
      info = node.get_info(show_all)
      for line in info:
        print line
    return 0

class TypeCmd(Command):
  def handle_vol(self, vol):
    p = self.opts
    if len(p) == 0:
      print "Usage: type <ami_file>"
      return 1
    else:
      name = make_fsstr(p[0])
      data = vol.read_file(name)
      print data
      return 0

class ReadCmd(Command):
  def handle_vol(self, vol):
    p = self.opts
    n = len(p)
    if n == 0 or n > 2:
      print "Usage: read <ami_file|dir> [sys_file]"
      return 1
    # determine output name
    out_name = os.path.basename(p[0])
    if n == 2:
      if os.path.isdir(p[1]):
        out_name = os.path.join(p[1],out_name)  
      else:
        out_name = p[1]
    # single file operation
    name = make_fsstr(p[0])
    node = vol.get_path_name(name)
    if node == None:
      print "Node not found:",p[0]
      return 2
    # its a file
    if node.is_file():
      data = node.get_file_data()
      # write data to file
      fh = open(out_name,"wb")
      fh.write(data)
      fh.close()
    # its a dir
    elif node.is_dir():
      img = Imager()
      img.unpack_dir(node, out_name)
    node.flush()
    return 0

class InfoCmd(Command):
  def handle_vol(self, vol):
    info = vol.get_info()
    for line in info:
      print line
    return 0

# ----- Edit Image -----

class MakeDirCmd(Command):
  def __init__(self, args, opts):
    Command.__init__(self, args, opts, edit=True)
  def handle_vol(self, vol):
    if len(self.opts) != 1:
      print "Usage: mkdir <dir_path>"
      return 1
    else:
      dir_path = make_fsstr(self.opts[0])
      vol.create_dir(dir_path)
      return 0

class WriteCmd(Command):
  def __init__(self, args, opts):
    Command.__init__(self, args, opts, edit=True)
  def handle_vol(self, vol):
    n = len(self.opts)
    if n == 0 or n > 2:
      print "Usage: write <sys_file|dir> [ami_path]"
      return 1
    # get file_name and ami_path
    sys_file = self.opts[0]
    file_name = os.path.basename(sys_file)
    if n > 1:
        ami_path = self.opts[1]
    else:
        ami_path = os.path.basename(sys_file)
    # check sys path
    if not os.path.exists(sys_file):
      print "File not found:",sys_file
      return 2
      
    ami_path = make_fsstr(ami_path)
    file_name = make_fsstr(file_name)
    # handle file
    if os.path.isfile(sys_file):
      fh = open(sys_file,"rb")
      data = fh.read()
      fh.close()
      vol.write_file(data, ami_path, file_name)
    # handle dir
    elif os.path.isdir(sys_file):
      parent_node, dir_name = vol.get_create_path_name(ami_path,file_name)
      if parent_node == None:
        print "Invalid path",ami_path
        return 2
      node = parent_node.create_dir(dir_name)
      img = Imager()
      img.pack_dir(sys_file, node)
      
    return 0
    
class DeleteCmd(Command):
  def __init__(self, args, opts):
    Command.__init__(self, args, opts, edit=True)
  def handle_vol(self, vol):
    n = len(self.opts)
    if n == 0:
      print "Usage: delete <ami_path> [wipe] [all]"
      return 1
    do_wipe = 'wipe' in self.opts
    do_all = 'all' in self.opts
    path = make_fsstr(self.opts[0])
    node = vol.delete(path, wipe=do_wipe, all=do_all)
    return 0
      
class ProtectCmd(Command):
  def __init__(self, args, opts):
    Command.__init__(self, args, opts, edit=True)
  def handle_vol(self, vol):
    n = len(self.opts)
    if n != 2:
      print "Usage: protect <ami_file> <protect>"
      return 1
    name = make_fsstr(self.opts[0])
    pr_str = self.opts[1]    
    node = vol.get_path_name(name)
    if node != None:
      node.change_protect_by_string(pr_str)
      return 0
    else:
      print "Can't find node:",name
      return 2

class CommentCmd(Command):
  def __init__(self, args, opts):
    Command.__init__(self, args, opts, edit=True)
  def handle_vol(self, vol):
    n = len(self.opts)
    if n != 2:
      print "Usage: comment <ami_file> <comment>"
      return 1
    name = make_fsstr(self.opts[0])
    comment = make_fsstr(self.opts[1])
    node = vol.get_path_name(name)
    if node != None:
      node.change_comment(comment)
      return 0
    else:
      print "Can't find node:",name
      return 2
    
class TimeCmd(Command):
  def __init__(self, args, opts):
    Command.__init__(self, args, opts, edit=True)
  def handle_vol(self, vol):
    n = len(self.opts)
    if n != 2:
      print "Usage: time <ami_file> <time>"
      return 1
    name = self.opts[0]
    tstr = self.opts[1]
    node = vol.get_path_name(name)
    if node != None:
      node.change_mod_ts_by_string(tstr)
      return 0
    else:
      print "Can't find node:",name
      return 2

class RelabelCmd(Command):
  def __init__(self, args, opts):
    Command.__init__(self, args, opts, edit=True)
  def handle_vol(self, vol):
    n = len(self.opts)
    if n != 1:
      print "Usage: relabel <new_name>"
      return 1
    name = self.opts[0]
    node = vol.relabel(name)
    return 0

# ----- Block Tools -----

class BlockCmd(Command):
  def handle_vol(self, vol):
    n = len(self.opts)
    if n == 0:
      print "Usage: block ( boot | root | node <ami_file> [data] | dump <block_no> )"
      return 1
    cmd = self.opts[0]
    if cmd == 'boot':
      vol.boot.dump()
      return 0
    elif cmd == 'root':
      vol.root.dump()
      return 0
    elif cmd == 'node':
      if n == 1:
        print "No node given!"
        return 1
      else:
        name = make_fsstr(self.opts[1])
        node = vol.get_path_name(name)
        if node != None:
          with_data = "data" in self.opts
          node.dump_blocks(with_data)
          return 0
        else:
          print "Can't find node:",name
          return 2
    elif cmd == 'dump':
      if n == 1:
        print "No block number given!"
        return 1
      else:
        block_no = int(self.opts[1])
        data = vol.blkdev.read_block(block_no)
        print_hex(data)

# ----- Bitmap Tools -----

class BitmapCmd(Command):
  def handle_vol(self, vol):
    n = len(self.opts)
    if n == 0:
      print "Usage: bitmap ( free | used | find [n] | all | maps | root [all] | node <path> [all] [entries]) [brief]"
      return 1
    cmd = self.opts[0]
    
    # brief mode
    brief = False
    if self.opts[-1] == 'brief':
      brief = True
      self.opts = self.opts[:-1]
  
    if cmd == 'free':
      vol.bitmap.print_free(brief)
      return 0
    elif cmd == 'used':
      vol.bitmap.print_used(brief)
      return 0
    elif cmd == 'find':
      if n == 2:
        num = int(self.opts[1])
        blk_nums = vol.bitmap.find_n_free(num)
        if blk_nums == None:
          print "No %d free blocks found" % num
          return 100
        else:
          print "Free %d blocks:" % num,blk_nums
          return 0
      else:
        blk_num = vol.bitmap.find_free()
        if blk_num == None:
          print "No free block found"
          return 100
        else:
          print "Free block:",blk_num
          return 0
    elif cmd == 'all':
      bm = vol.bitmap.create_draw_bitmap()
      vol.bitmap.draw_on_bitmap(bm)
      vol.root_dir.draw_on_bitmap(bm, True)
      vol.bitmap.print_draw_bitmap(bm, brief)
      return 0
    elif cmd == 'maps':
      bm = vol.bitmap.create_draw_bitmap()
      vol.bitmap.draw_on_bitmap(bm)
      vol.bitmap.print_draw_bitmap(bm, brief)
      return 0
    elif cmd == 'root':
      show_entries = 'entries' in self.opts
      bm = vol.bitmap.create_draw_bitmap()
      vol.root_dir.draw_on_bitmap(bm, False, show_entries)
      vol.bitmap.print_draw_bitmap(bm, brief)
      return 0
    elif cmd == 'node':
      if n > 1:
        name = make_fsstr(self.opts[1])
        node = vol.get_path_name(name)
        if node != None:
          show_all = 'all' in self.opts
          show_entries = 'entries' in self.opts
          bm = vol.bitmap.create_draw_bitmap()
          node.draw_on_bitmap(bm, show_all, show_entries)
          vol.bitmap.print_draw_bitmap(bm, brief)
          return 0
        else:
          print "Node '%s' not found!" % self.opts[1]
          return 2
      else:
        print "Need node path!"
        return 1
    else:
      print "Unknown bitmap command!"
      return 1

# ----- Root Tools -----

class RootCmd(Command):
  def handle_vol(self, volume):
    n = len(self.opts)
    if n == 0:
      print "Usage: root ( show | create_time <time> | disk_time <time> | time <time> )"
      return 1
    cmd = self.opts[0]
    # root show
    if cmd == 'show':
      volume.root.dump()
      return 0
    # create_time <time>
    elif cmd == 'create_time':
      if n != 2:
        print "create_time <time>"
        return 1
      else:
        volume.change_create_ts_by_string(self.opts[1])
        return 0
    # disk_time <time>
    elif cmd == 'disk_time':
      if n != 2:
        print "disk_time <time>"
        return 1
      else:
        volume.change_disk_ts_by_string(self.opts[1])
        return 0
    # time <time>
    elif cmd == 'time':
      if n != 2:
        print "time <time>"
        return 1
      else:
        volume.change_mod_ts_by_string(self.opts[1])
        return 0
    # boot ?
    else:
      print "Unknown root command!"
      return 1

# ----- Boot Tools -----

class BootCmd(Command):
  def handle_blkdev(self, blkdev):
    n = len(self.opts)
    if n == 0:
      print "Usage: boot ( show [hex] [asm] | read <file> | write <file> | install [boot1x] | clear )"
      return 1
    cmd = self.opts[0]
    # fetch boot blcok
    bb = BootBlock(blkdev, 0)
    bb.read()
    # boot show [hex] [asm]
    if cmd == 'show':
      bb.dump()
      if bb.valid and bb.boot_code != None:
        if 'hex' in self.opts:
          print_hex(bb.boot_code)
        if 'asm' in self.opts:
          dis = DisAsm()
          code = dis.disassemble(bb.boot_code)
          dis.dump(code)
      return 0
    # boot read <file>
    elif cmd == 'read':
      if n > 1:
        if bb.valid:
          if bb.boot_code != None:
            f = open(self.opts[1], "wb")
            f.write(bb.boot_code)
            f.close()
            return 0
          else:
            print "No Boot Code found!"
            return 2
        else:
          print "Invalid Boot Block!"
          return 1
      else:
        print "No file name!"
        return 1
    # boot write <file>
    elif cmd == 'write':
      if n > 1:
        f = open(self.opts[1], "rb")
        data = f.read()
        f.close()
        ok = bb.set_boot_code(data)
        if ok:
          bb.write()
          return 0
        else:
          print "Boot Code invalid!"
          return 2
      else:
        print "No file name!"
        return 1
    # boot install [<name>]
    elif cmd == 'install':
      if n == 1: # default boot code
        name = "boot2x3x"
      else: 
        name = self.opts[1]
      # boot code directory
      bc_dir = bb.get_boot_code_dir()
      if bc_dir == None:
        print "No boot code directory found!"
        return 1
      path = os.path.join(bc_dir, name + ".bin")
      if not os.path.exists:
        print "Boot code '%s' not found!" % path
        return 1
      # read data
      f = open(path,"rb")
      data = f.read()
      f.close()
      ok = bb.set_boot_code(data)
      if ok:
        bb.write()
        return 0
      else:
        print "Boot Code invalid!"
        return 2
    # boot clear
    elif cmd == 'clear':
      bb.set_boot_code(None)
      bb.write()
      return 0
    # boot ?
    else:
      print "Unknown boot command!"
      return 1

# ----- BlkDev Command -----

class BlkDevCmd(Command):
  def handle_blkdev(self, blkdev):
    blkdev.dump()
    return 0

# ----- main -----
# call scanner and process all files with selected command
cmd_map = {
"open" : OpenCmd,
"create" : CreateCmd,
"list" : ListCmd,
"type" : TypeCmd,
"read" : ReadCmd,
"makedir" : MakeDirCmd,
"write" : WriteCmd,
"delete" : DeleteCmd, 
"format" : FormatCmd,
"bitmap" : BitmapCmd,
"blkdev" : BlkDevCmd,
"protect" : ProtectCmd,
"comment" : CommentCmd,
"time" : TimeCmd,
"block" : BlockCmd,
"pack" : PackCmd,
"unpack" : UnpackCmd,
"repack" : RepackCmd,
"boot" : BootCmd,
"root" : RootCmd,
"info" : InfoCmd,
"relabel" : RelabelCmd
}

parser = argparse.ArgumentParser()
parser.add_argument('image_file')
parser.add_argument('command_list', nargs='+', help="command: "+",".join(cmd_map.keys()))
parser.add_argument('-v', '--verbose', action='store_true', default=False, help="be more verbos")
parser.add_argument('-s', '--seperator', default='+', help="set the command separator char sequence")
parser.add_argument('-r', '--read-only', action='store_true', default=False, help="read-only operation")
parser.add_argument('-f', '--force', action='store_true', default=False, help="force overwrite existing image")
args = parser.parse_args()

cmd_list = args.command_list
sep = args.seperator
queue = FSCommandQueue(args, cmd_list, sep, cmd_map)
code = queue.run()
sys.exit(code)