#!/usr/bin/python

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)

cwd = '/'
jolietcwd = '/'
rock_ridge_names = False
joliet_names = False
class PyCdlibCmdLoop(cmd.Cmd):
    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_rock_ridge_names(self):
        print("> rock_ridge_names")
        print("Toggle whether to print Rock Ridge names or not.  If Rock Ridge names are being printed, Joliet names are automatically disabled")

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

        global rock_ridge_names
        global joliet_names

        if not iso.rock_ridge:
            print("Can only enable Rock Ridge names for Rock Ridge ISOs")
            return False

        if rock_ridge_names:
            rock_ridge_names = False
        else:
            joliet_names = False
            rock_ridge_names = True

    def help_joliet_names(self):
        print("> joliet_names")
        print("Toggle whether to print Joliet names or not.  If Joliet names are being printed, Rock Ridge names are automatically disabled")

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

        global rock_ridge_names
        global joliet_names

        if iso.joliet_vd is None:
            print("Can only enable Joliet names for Joliet ISOs")
            return False

        if joliet_names:
            joliet_names = False
        else:
            rock_ridge_names = False
            joliet_names = True

    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

        global cwd
        global jolietcwd
        global rock_ridge_names
        global joliet_names

        working_dir = cwd
        if joliet_names:
            working_dir = jolietcwd

        for child in iso.list_dir(working_dir, joliet_names):
            prefix = "F"
            if child.is_dir():
                prefix = "D"

            name = child.file_identifier()
            if rock_ridge_names:
                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))

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

    def do_cd(self, line):
        global jolietcwd
        global cwd

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

        directory = split[0]

        if directory == '/':
            tmp = '/'
        else:
            tmp = os.path.normpath(os.path.join(cwd, directory))

        rec = iso.get_entry(tmp, joliet_names)
        if not rec.is_dir():
            print("Entry %s is not a directory" % (directory))
            return

        if joliet_names:
            jolietcwd = tmp
        else:
            cwd = tmp

    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

        global cwd
        global jolietcwd

        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 joliet_names:
                iso_file = os.path.join(jolietcwd, iso_file)
            else:
                iso_file = os.path.join(cwd, iso_file)

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

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

    def do_cwd(self, line):
        global cwd
        global jolietcwd

        if len(line) != 0:
            print("No parameters allowed for cwd")
            return

        if joliet_names:
            print(jolietcwd)
        else:
            print(cwd)

    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):
        global cwd
        global jolietcwd

        if len(line) != 0:
            print("No parameters allowed for tree")
            return

        if joliet_names:
            working_dir = jolietcwd
        else:
            working_dir = 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(working_dir, joliet_names)
        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))

    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

        out_name = split[0]

        iso.write(out_name)

    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

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

        if iso_path[0] != '/':
            iso_path = os.path.join(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)

    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

        iso_path = split[0]

        if iso_path[0] != '/':
            iso_path = os.path.join(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)

    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

        iso_path = split[0]

        if iso_path[0] != '/':
            iso_path = os.path.join(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)

    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

        iso_path = split[0]

        if iso_path[0] != '/':
            iso_path = os.path.join(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)

    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

        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()

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

iso.close()
fp.close()
