#!/usr/bin/python3 -s

import os
import sys
import cmd
import collections

import networkx
import matplotlib.pyplot

import pycdlib

def usage():
    print("Usage: %s <isofile>" % (sys.argv[0]))

if len(sys.argv) != 2:
    usage()
    sys.exit(1)

iso = pycdlib.PyCdlib()
fp = open(sys.argv[1], 'r')
iso.open_fp(fp)

class PyCdlibCmdLoop(cmd.Cmd):
    def __init__(self):
        cmd.Cmd.__init__(self)
        self.cwd = '/'
        self.jolietcwd = '/'
        self.print_mode = 'iso9660'

    prompt = '(pycdlib) '

    def help_exit(self):
        print("> exit")
        print("Exit the program.")

    def do_exit(self, line):
        if len(line) != 0:
            print("No parameters allowed for exit")
            return False
        return True

    def help_quit(self):
        print("> quit")
        print("Exit the program.")

    def do_quit(self, line):
        if len(line) != 0:
            print("No parameters allowed for quit")
            return False
        return True

    def do_EOF(self, line):
        print
        return True

    def help_print_mode(self):
        print("> print_mode [iso9660|rr|joliet]")
        print("Change which 'mode' of filenames are printed out.  There are three main\n"
              "modes: ISO9660 (iso9660, the default), rr (Rock Ridge extensions),\n"
              "and Joliet (joliet).  The original iso9660 mode only allows filenames\n"
              "of 8 characters, plus 3 for the extension.  The Rock Ridge extensions\n"
              "allow much deeper filenames and directory structures.  The Joliet\n"
              "extensions also allow much deeper filenames and directory structures, but\n"
              "in an entirely different namespace (though in most circumstances, the\n"
              "Joliet namespace will mirror the ISO9660/Rock Ridge namespace).  Note that\n"
              "any given ISO will always have ISO9660 mode, but may have any combination\n"
              "of Rock Ridge and Joliet (including neither).  Running this command with\n"
              "no arguments prints out the current mode.  Passing 'iso9660' as an\n"
              "argument sets it to the original ISO9660 mode.  Passing 'rr' as an\n"
              "argument sets it to Rock Ridge mode.  And passing 'joliet' as an argument\n"
              "sets it to the Joliet mode.")

    def do_print_mode(self, line):
        split = line.split()
        if len(split) == 0:
            print(self.print_mode)
            return False
        elif len(split) != 1:
            print("Only a single parameter allowed for print_mode")
            return False

        if split[0] not in ['iso9660', 'rr', 'joliet']:
            print("Parameter for print_mode must be one of 'iso9660', 'rr', or 'joliet'")
            return False

        if split[0] == 'rr' and not iso.rock_ridge:
            print("Can only enable Rock Ridge names for Rock Ridge ISOs")
            return False

        if split[0] == 'joliet' and iso.joliet_vd is None:
            print("Can only enable Joliet names for Joliet ISOs")
            return False

        self.print_mode = split[0]

        return False

    def help_ls(self):
        print("> ls")
        print("Show the contents of the current working directory. The format of the output is:\n")
        print("TYPE(F=file, D=directory) NAME")

    def do_ls(self, line):
        if len(line) != 0:
            print("No parameters allowed for ls")
            return

        wd = self.jolietcwd if self.print_mode == 'joliet' else self.cwd

        for child in iso.list_dir(wd, self.print_mode == 'joliet'):
            prefix = "F"
            if child.is_dir():
                prefix = "D"

            name = child.file_identifier()
            if self.print_mode == 'rr':
                name = ""
                if child.is_dot():
                    name = "."
                elif child.is_dotdot():
                    name = ".."
                else:
                    if child.rock_ridge is not None:
                        if child.rock_ridge.name() != "":
                            name = "%s" % (child.rock_ridge.name())
                            if child.rock_ridge.is_symlink():
                                name += " -> %s" % (child.rock_ridge.symlink_path())
                                prefix += "S"

            print("%2s %s" % (prefix, name))

        return False

    def help_cd(self):
        print("> cd <iso_dir>")
        print("Change directory to <iso_dir> on the ISO.")

    def do_cd(self, line):
        split = line.split()
        if len(split) != 1:
            print("The cd command supports one and only one parameter")
            return False

        directory = split[0]

        if directory == '/':
            tmp = '/'
        else:
            wd = self.jolietcwd if self.print_mode == 'joliet' else self.cwd
            tmp = os.path.normpath(os.path.join(wd, directory))

        rec = iso.get_entry(tmp, self.print_mode == 'joliet')
        if not rec.is_dir():
            print("Entry %s is not a directory" % (directory))
            return False

        if self.print_mode == 'joliet':
            self.jolietcwd = tmp
        else:
            self.cwd = tmp

        return False

    def help_get(self):
        print("> get <iso_file> <out_file>")
        print("Get the contents of <iso_file> from the ISO and write them to <out_file>.")

    def do_get(self, line):
        split = line.split()
        if len(split) != 2:
            print("The get command must be passed two parameters.")
            return False

        iso_file = split[0]
        outfile = split[1]

        if iso_file[0] != '/':
            # In this case, it is a relative path, so we should prepend the
            # cwd.
            if self.print_mode == 'joliet':
                iso_file = os.path.join(self.jolietcwd, iso_file)
            else:
                iso_file = os.path.join(self.cwd, iso_file)

        with open(outfile, 'wb') as outfp:
            iso.get_and_write_fp(iso_file, outfp)

        return False

    def help_cwd(self):
        print("> cwd")
        print("Show the current working directory.")

    def do_cwd(self, line):
        if len(line) != 0:
            print("No parameters allowed for cwd")
            return False

        if self.print_mode == 'joliet':
            print(self.jolietcwd)
        else:
            print(self.cwd)

        return False

    def help_tree(self):
        print("> tree")
        print("Prints all files and subdirectories below the current directory (similar to the Unix 'tree' command).")

    def do_tree(self, line):
        if len(line) != 0:
            print("No parameters allowed for tree")
            return False

        wd = self.jolietcwd if self.print_mode == 'joliet' else self.cwd

        utf8_copyright = "\302\251"
        utf8_corner = "\342\224\224\342\224\200\342\224\200"
        utf8_middlebar = "\342\224\234\342\224\200\342\224\200"
        utf8_vertical_line = "\342\224\202\302\240\302\240"
        entry = iso.get_entry(wd, self.print_mode == 'joliet')
        dirs = collections.deque([(entry, 0, [])])
        while dirs:
            dir_record,depth,lasts = dirs.popleft()
            prefix = ""
            for index,last in enumerate(lasts):
                if last:
                    if index == (len(lasts) - 1):
                        prefix += utf8_corner + " "
                    else:
                        prefix += "    "
                else:
                    if index == (len(lasts) - 1):
                        prefix += utf8_middlebar + " "
                    else:
                        prefix += utf8_vertical_line + " "
            print("%s%s" % (prefix, dir_record.file_identifier()))

            if dir_record.is_dir():
                tmp = []
                for index,child in enumerate(dir_record.children):
                    if child.is_dot() or child.is_dotdot():
                        continue
                    last = False
                    if index == (len(dir_record.children) - 1):
                        last = True
                    tmp.append((child, depth+1, lasts + [last]))
                dirs.extendleft(reversed(tmp))

        return False

    def help_write(self):
        print("> write <out_file>")
        print("Write the current ISO contents to <out_file>.")

    def do_write(self, line):
        split = line.split()
        if len(split) != 1:
            print("The write command supports one and only one parameter.")
            return False

        out_name = split[0]

        iso.write(out_name)

        return False

    def help_add_file(self):
        print("> add_file <iso_path> <src_filename>")
        print("Add the contents of <src_filename> to the ISO at the location specified in <iso_path>.")

    def do_add_file(self, line):
        split = line.split()
        if len(split) != 2:
            print("The add_file command needs two parameters (iso path and source file path).")
            return False

        iso_path = split[0]
        src_path = split[1]

        if iso_path[0] != '/':
            iso_path = os.path.join(self.cwd, iso_path)

        rrname = None
        if iso.rock_ridge:
            rrname = os.path.basename(iso_path)

        jpath = None
        if iso.joliet_vd is not None:
            jpath = iso_path

        iso.add_file(src_path, iso_path, rr_name=rrname, joliet_path=jpath)

        return False

    def help_rm_file(self):
        print("> rm_file <iso_path>")
        print("Remove the contents of <iso_path> from the ISO.")

    def do_rm_file(self, line):
        split = line.split()
        if len(split) != 1:
            print("The rm_file command takes one and only one parameter (iso path).")
            return False

        iso_path = split[0]

        if iso_path[0] != '/':
            iso_path = os.path.join(self.cwd, iso_path)

        rrname = None
        if iso.rock_ridge:
            rrname = os.path.basename(iso_path)

        jpath = None
        if iso.joliet_vd is not None:
            jpath = iso_path

        iso.rm_file(iso_path, rr_name=rrname, joliet_path=jpath)

        return False

    def help_mkdir(self):
        print("> mkdir <iso_path>")
        print("Make a new directory called <iso_path>.")

    def do_mkdir(self, line):
        split = line.split()
        if len(split) != 1:
            print("The mkdir command takes one and only one parameter (iso path).")
            return False

        iso_path = split[0]

        if iso_path[0] != '/':
            iso_path = os.path.join(self.cwd, iso_path)

        rrname = None
        if iso.rock_ridge:
            rrname = os.path.basename(iso_path)

        jpath = None
        if iso.joliet_vd is not None:
            jpath = iso_path

        iso.add_directory(iso_path, rr_name=rrname, joliet_path=jpath)

        return False

    def help_rmdir(self):
        print("> rmdir <iso_path>")
        print("Remove the directory at <iso_path>.  Note that the directory must be empty for the command to succeed.")

    def do_rmdir(self, line):
        split = line.split()
        if len(split) != 1:
            print("The rmdir command takes one and only one parameter (iso path).")
            return False

        iso_path = split[0]

        if iso_path[0] != '/':
            iso_path = os.path.join(self.cwd, iso_path)

        rrname = None
        if iso.rock_ridge:
            rrname = os.path.basename(iso_path)

        jpath = None
        if iso.joliet_vd is not None:
            jpath = iso_path

        iso.rm_directory(iso_path, rr_name=rrname, joliet_path=jpath)

        return False

    def help_graph(self):
        print("> graph")
        print("Print a graph representing the layout of the ISO.")

    def do_graph(self, line):
        if len(line) != 0:
            print("No parameters allowed for graph")
            return False

        iso9660_graph = networkx.DiGraph()
        rr_graph = networkx.DiGraph()
        dirs = collections.deque([('/', '/')])
        while dirs:
            curr,rr_curr = dirs.popleft()
            for child in iso.list_dir(curr):
                if child.is_dot() or child.is_dotdot():
                    continue
                parent_ident = curr
                if parent_ident != '/':
                    parent_ident = parent_ident.split('/')[-1]
                iso9660_graph.add_edge(parent_ident, child.file_identifier())
                if child.rock_ridge is not None:
                    rr_graph.add_edge(rr_curr, child.rock_ridge.name())
                if child.is_dir():
                    append = curr
                    rr_append = ""
                    if append[-1] != '/':
                        append += '/'
                    rr_name = ""
                    if child.rock_ridge is not None:
                        rr_append = rr_curr
                        if rr_append[-1] != '/':
                            rr_append += '/'
                        rr_name = rr_append + child.rock_ridge.name()
                    dirs.append((append + child.file_identifier(), rr_name))

        isofig = matplotlib.pyplot.figure()
        isofig.suptitle('ISO9660')
        isofig.canvas.set_window_title('ISO9660')
        pos = networkx.graphviz_layout(iso9660_graph, prog='dot')
        networkx.draw_networkx(iso9660_graph, pos=pos, arrows=True)
        if len(rr_graph) > 0:
            rrfig = matplotlib.pyplot.figure()
            rrfig.suptitle('Rock Ridge')
            rrfig.canvas.set_window_title('Rock Ridge')
            pos = networkx.graphviz_layout(rr_graph, prog='dot')
            networkx.draw_networkx(rr_graph, pos=pos, arrows=True)
        matplotlib.pyplot.show()

        return False

done = False
while not done:
    try:
        PyCdlibCmdLoop().cmdloop()
        done = True
    except Exception as e:
        print e

iso.close()
fp.close()
