path: root/tools/fsmaker/sh.py

import os
import sys
import cmd
import api

import math
import stat
import errno
import shlex
import Queue
import struct
import string
import optparse
import tempfile

import curses.ascii

class OptionParser(optparse.OptionParser):

    def __init__(self, **args):
        optparse.OptionParser.__init__(self, **args)

    def exit(self, code=0, msg=""):
        sys.stderr.write(msg)

class FsmakerShell(cmd.Cmd):

    def __init__(self, simdisk):
        self._simdisk = simdisk
        self._curdir = "/"
        self.prompt = "{0} > ".format(self._curdir)
        cmd.Cmd.__init__(self)

        self._parse_superblock = OptionParser(usage="usage: %prog", prog="superblock", description="prints a summary of the superblock's contents")

        self._parse_inode = OptionParser(usage="usage: %prog <nums...>", prog="inode", description="prints a summary of the specified inode's contents")
        self._parse_inode.add_option("-i", "--indirect", action="store_true", default=False,
                                     help="if the inode is a directory or data file print the indirect block contents")
        self._parse_inode.add_option("-c", "--contents", action="store_true", default=False,
                                     help="if the inode is a data file or directory this prints the contents of the file as part of the summary")
        self._parse_inode.add_option("-l", "--list", action="store_true", default=False,
                                     help="if the inode is a directory this prints a directory listing as part of the summary")

        self._parse_block = OptionParser(usage="usage: %prog", prog="block", description="prints the data from a given block")
        self._parse_ls = OptionParser(usage="usage: %prog <dirs...>", prog="ls", description="prints a directory listing")
        self._parse_cat = OptionParser(usage="usage: %prog <files...>", prog="cat", description="prints the contents of a file")
        self._parse_cd = OptionParser(usage="usage: %prog <dir>", prog="cd", description="changes the current working directory")

        self._parse_trunc = OptionParser(usage="usage: %prog <files...>", prog="truncate", description="changes the size of a file")
        self._parse_trunc.add_option("-s", "--size", action="store", type="int", default=0,
                                     help="the new file size (defaults to %default)")

        self._parse_rm = OptionParser(usage="usage: %prog <files...>", prog="rm", description="removes a file from the disk")
        self._parse_rm.add_option("-r", "--recursive", action="store_true", default=False,
                                     help="recrsively destroy all subdirectories and files of any directory arguments")
        self._parse_rmdir = OptionParser(usage="usage: %prog <dirs...>", prog="rmdir", description="removes an empty directory from the disk")

        self._parse_touch = OptionParser(usage="usage: %prog <dirs...>", prog="touch", description="creates a plain data file")
        self._parse_mkdir = OptionParser(usage="usage: %prog <dirs...>", prog="mkdir", description="creates an empty directory")

        self._parse_getfile = OptionParser(usage="usage: %prog <source> <dest>", prog="getfile", description="gets a file from the real disk and puts it on the simdisk")
        self._parse_putfile = OptionParser(usage="usage: %prog <source> <dest>", prog="putfile", description="puts a file from the simdisk onto the real disk")

        self._parse_format = OptionParser(usage="usage: %prog -i <inode count> [-s <size>|-b <blocks>]", prog="format", description="formats the simdisk to an empty file system")
        self._parse_format.add_option("-s", "--size", action="store", type="int", default=None,
                                      help="size for the new file system in bytes, must specify either this option or -b but not both")
        self._parse_format.add_option("-b", "--blocks", action="store", type="int", default=None,
                                      help="size for the new file system in blocks, must specify either this option or -s but not both")
        self._parse_format.add_option("-i", "--inodes", action="store", type="int", default=None,
                                      help="number of inodes to put on the disk, this must be specified and be compatible with the size of the disk (there must be enough space for the inodes)")
        self._parse_format.add_option("-d", "--directory", action="store", type="str", default=None,
                                      help="initializes the disk with the contents of the specified directory")

    def open(self, path, create=False):
        if (path.startswith("/")):
            return self._simdisk.open(path, create=create)
        else:
            return self._simdisk.open(self._curdir + "/" + path, create=create)

    def filepath_completion(self, text, line, begin, end, types=api.S5_TYPES):
        res = []
        filename = text

        bline = line[:begin]
        try:
            parts = shlex.split(bline)
            if (bline.endswith('"')):
                return []
            if (bline.endswith(" ") and not bline.endswith("\\ ")):
                dirpath = self._curdir
            else:
                dirpath = parts[-1]
        except ValueError:
            dirpath = bline.rsplit('"', 1)[-1]
            if (len(dirpath.strip()) == 0):
                dirpath = self._curdir

        try:
            dirinode = self.open(dirpath)
            if (None != dirinode):
                for dirent in dirinode.getdents():
                    if (dirent.name.startswith(filename) and dirent.name != "." and dirent.name != ".."):
                        t = self._simdisk.get_inode(dirent.inode).get_type()
                        if (t in types or t == api.S5_TYPE_DIR):
                            if (t == api.S5_TYPE_DIR):
                                res.append(dirent.name + "/")
                            else:
                                res.append(dirent.name)
                return res
            else:
                return []
        except api.S5fsException as e:
            return []

    def real_filepath_completion(self, text, line, begin, end):
        res = []
        filename = text

        bline = line[:begin]
        try:
            parts = shlex.split(bline)
            if (bline.endswith('"')):
                return []
            if (bline.endswith(" ") and not bline.endswith("\\ ")):
                dirpath = os.curdir
            else:
                dirpath = parts[-1]
        except ValueError:
            dirpath = bline.rsplit('"', 1)[-1]
            if (len(dirpath.strip()) == 0):
                dirpath = os.curdir

        try:
            for dirent in os.listdir(dirpath):
                if (dirent.startswith(filename)):
                    if (os.path.isdir(os.path.join(dirpath, dirent))):
                        res.append(dirent + "/")
                    else:
                        res.append(dirent)
            return res
        except api.S5fsException as e:
            return []

    def completion_argnum(self, text, line, begin, end):
        bline = line[:begin]
        try:
            parts = [x for x in shlex.split(bline) if not x.startswith("-")]
            if (bline.endswith(" ") and not bline.endswith("\\ ")):
                return len(parts)
            else:
                return len(parts) - 1
        except ValueError:
            parts = [x for x in shlex.split(bline.rsplit('"', 1)[0]) if not x.startswith("-")]
            return len(parts)

    def get_parentdir(self, path):
        path = os.path.normpath(path).rsplit('/', 1)
        if (len(path) == 1):
            res = (self.open(self._curdir), path[0])
        else:
            if (len(path[1]) == 0):
                res = (self.open(path[0]), '.')
            else:
                res = (self.open(path[0]), path[1])
        if (res[0] == None):
            raise api.S5fsException("no such directory: {0}".format(path[0]))
        return res

    def binary_print(self, data, prefix=""):
        binary = ""
        line = ""
        for i in xrange(int(math.floor(len(data) / 2))):
            unpacked = struct.unpack("BB", data[i * 2:(i + 1) * 2])
            binary += "{0:02x}{1:02x} ".format(*unpacked)
            for d in unpacked:
                if (curses.ascii.isprint(d)):
                    line += chr(d)
                else:
                    line += '.'
            if ((i + 1) % 10 == 0):
                print("{2}{0:<50} {1}".format(binary, line, prefix))
                binary = ""
                line = ""
        if (len(data) % 2 != 0):
            d = struct.unpack("B", data[-1])[0]
            binary += "{0:02x}".format(d)
            if (curses.ascii.isprint(d)):
                line += chr(d)
            else:
                line += '.'
        if (binary != ""):
            print("{2}{0:<50} {1}".format(binary, line, prefix))

    def dirents_print(self, dirents, prefix=""):
        direntlist = []
        maxlen = 0
        maxinode = 0
        for dirent in dirents:
            maxlen = max(maxlen, len(dirent.name))
            maxinode = max(maxinode, dirent.inode)
            direntlist.append(dirent)
        if (len(direntlist) == 0):
            print ("{0}<empty directory>".format(prefix))
        else:
            for dirent in direntlist:
                try:
                    inode = self._simdisk.get_inode(dirent.inode)
                    itype = inode.get_type()
                    itypestr = inode.get_type_str(short=True)
                    isize = inode.get_size()
                except api.S5fsException as e:
                    inode = None
                    errmsg = str(e)
                if (inode == None):
                    msg = "<{0}>".format(errmsg)
                else:
                    if (itype in set([ api.S5_TYPE_DATA, api.S5_TYPE_DIR ])):
                        msg = "{0} {1} bytes".format(itypestr, isize)
                    elif (itype in set([ api.S5_TYPE_BLK, api.S5_TYPE_CHR ])):
                        msg = "{0}".format(itypestr)
                    else:
                        msg = "{0} (INVALID, free inode)".format(itypestr)
                print("{4}{0:<{1}} {2:>{3}} {5}".format(dirent.name, maxlen, dirent.inode, len(str(maxinode)), prefix, msg))

    def help_help(self):
        print("prints help information about a command")

    def do_cd(self, args):
        try:
            (options, args) = self._parse_cd.parse_args(shlex.split(args))
        except ValueError as e:
            self._parse_cd.error(str(e))
            return

        if (len(args) != 1):
            self._parse_cd.error("command takes exactly one argument")
        else:
            f = self.open(args[0])
            if (None == f):
                self._parse_cd.error("no such file or directory: {0}".format(args[0]))
            elif (api.S5_TYPE_DIR != f.get_type()):
                self._parse_cd.error("not a directory: {0}".format(args[0]))
            else:
                self._curdir = os.path.normpath(self._curdir + ("" if self._curdir.endswith("/") else "/") + args[0])
                self.prompt = "{0} > ".format(self._curdir)

    def help_cd(self):
        self._parse_cd.print_help()

    def complete_cd(self, text, line, begidx, endidx):
        if (self.completion_argnum(text, line, begidx, endidx) == 1):
            return self.filepath_completion(text, line, begidx, endidx, types=set([ api.S5_TYPE_DIR ]))
        else:
            return []

    def do_superblock(self, args):
        try:
            (options, args) = self._parse_superblock.parse_args(shlex.split(args))
        except ValueError as e:
            self._parse_superblock.error(str(e))
            return

        if (len(args) != 0):
            self._parse_superblock.error("command does not take arguments")
        else:
            print(self._simdisk.get_super_block_summary())

    def help_superblock(self):
        self._parse_superblock.print_help()

    def complete_superblock(self, text, line, begidx, endidx):
        return []

    def do_inode(self, args):
        try:
            (options, args) = self._parse_inode.parse_args(shlex.split(args))
        except ValueError as e:
            self._parse_inode.error(str(e))
            return

        if (len(args) == 0):
            self._parse_inode.error("command requires at least one argument")
        else:
            for arg in args:
                try:
                    inode = self._simdisk.get_inode(string.atoi(arg))
                except ValueError as e:
                    inode = self.open(arg)
                except api.S5fsException as e:
                    inode = self.open(arg)

                if (None == inode):
                    self._parse_inode.error(arg + " is not a valid inode number or path")
                else:
                    try:
                        print(inode.get_summary())
                        if (options.indirect and inode.get_type() in set([ api.S5_TYPE_DATA, api.S5_TYPE_DIR ]) and inode.get_indirect_blockno() != 0):
                            try:
                                iblock = self._simdisk.get_block(inode.get_indirect_blockno())
                                for i in xrange(api.S5_BLOCK_SIZE / 4):
                                    num = struct.unpack("I", iblock.read(i * 4, 4))[0]
                                    sys.stdout.write(" {0:5}".format(num))
                                    if ((i + 1) % 8 == 0):
                                        sys.stdout.write("\n")
                                if ((i + 1) % 8 != 0):
                                    sys.stdout.write("\n")
                            except api.S5fsException as e:
                                self._parse_inode.error(str(e))
                        if (options.contents and inode.get_type() in set([ api.S5_TYPE_DATA, api.S5_TYPE_DIR ])):
                            print("contents:")
                            self.binary_print(inode.read(),prefix="  ")
                        if (options.list and inode.get_type() == api.S5_TYPE_DIR):
                            print("directory listing:")
                            self.dirents_print(inode.getdents(),prefix="  ")
                    except api.S5fsException as e:
                        self._parse_inode.error(str(e))
                print("- - - - -")

    def help_inode(self):
        self._parse_inode.print_help()

    def complete_inode(self, text, line, begidx, endidx):
        res = []
        try:
            string.atoi(text)
            res += [str(x) for x in xrange(self._simdisk.get_num_inodes()) if str(x).startswith(text)]
        except ValueError:
            None
        res += self.filepath_completion(text, line, begidx, endidx)
        return res

    def do_block(self, args):
        try:
            (options, args) = self._parse_block.parse_args(shlex.split(args))
        except ValueError as e:
            self._parse_block.error(str(e))
            return

        if (len(args) == 0):
            self._parse_block.error("command requires at least one argument")
        else:
            for arg in args:
                try:
                    blockno = string.atoi(arg)
                    print("block {0}:".format(blockno))
                    block = self._simdisk.get_block(blockno)
                    self.binary_print(block.read())
                except api.S5fsException as e:
                    self._parse_block.error(str(e))
                except ValueError as e:
                    self._parse_block.error(arg + " is not a valid block number")
                print("- - - - -")

    def help_block(self):
        self._parse_block.print_help()

    def do_ls(self, args):
        try:
            (options, args) = self._parse_ls.parse_args(shlex.split(args))
        except ValueError as e:
            self._parse_ls.error(str(e))
            return

        for arg in (args if len(args) > 0 else [self._curdir]):
            try:
                f = self.open(arg)
                if (f == None):
                    self._parse_ls.error("no such file or directory: {0}".format(arg))
                elif (f.get_type() != api.S5_TYPE_DIR):
                    self._parse_ls.error("not a directory: {0}".format(arg))
                else:
                    self.dirents_print(f.getdents())
            except api.S5fsException as e:
                self._parse_ls.error(str(e))
            print("- - - - -")

    def help_ls(self):
        self._parse_ls.print_help()

    def complete_ls(self, text, line, begidx, endidx):
        return self.filepath_completion(text, line, begidx, endidx, types=set([ api.S5_TYPE_DIR ]))

    def do_cat(self, args):
        try:
            (options, args) = self._parse_cat.parse_args(shlex.split(args))
        except ValueError as e:
            self._parse_cat.error(str(e))
            return

        if (len(args) == 0):
            self._parse_cat.error("command requires at least one argument")
        else:
            for arg in args:
                try:
                    f = self.open(arg)
                    if (f == None):
                        self._parse_cat.error("no such file or directory: {0}".format(arg))
                    elif (f.get_type() != api.S5_TYPE_DATA):
                        self._parse_cat.error("not a plain data file: {0}".format(arg))
                    else:
                        print(f.read())
                except api.S5fsException as e:
                    self._parse_cat.error(str(e))
            print("- - - - -")

    def help_cat(self):
        self._parse_cat.print_help()

    def complete_cat(self, text, line, begidx, endidx):
        return self.filepath_completion(text, line, begidx, endidx, types=set([ api.S5_TYPE_DATA ]))

    def do_truncate(self, args):
        try:
            (options, args) = self._parse_trunc.parse_args(shlex.split(args))
        except ValueError as e:
            self._parse_trunc.error(str(e))
            return

        if (len(args) == 0):
            self._parse_trunc.error("command requires at least one argument")
        else:
            for arg in args:
                try:
                    f = self.open(arg)
                    if (f == None):
                        self._parse_trunc.error("no such file or directory: {0}".format(arg))
                    elif (f.get_type() != api.S5_TYPE_DATA):
                        self._parse_trunc.error("not a plain data file: {0}".format(arg))
                    else:
                        f.truncate(size=options.size)
                except api.S5fsException as e:
                    self._parse_trunc.error(str(e))

    def help_truncate(self):
        self._parse_trunc.print_help()

    def complete_truncate(self, text, line, begidx, endidx):
        return self.filepath_completion(text, line, begidx, endidx, types=set([ api.S5_TYPE_DATA ]))

    def _dir_clear(self, directory):
        for dirent in directory.getdents():
            if (dirent.name != "." and dirent.name != ".."):
                inode = self._simdisk.get_inode(dirent.inode)
                if (inode.get_type() == api.S5_TYPE_DIR):
                    self._dir_clear(inode)
                    directory.rmdir(dirent.name)
                else:
                    directory.unlink(dirent.name)

    def do_rm(self, args):
        try:
            (options, args) = self._parse_rm.parse_args(shlex.split(args))
        except ValueError as e:
            self._parse_rm.error(str(e))
            return

        if (len(args) == 0):
            self._parse_rm.error("command requires at least one argument")
        else:
            for arg in args:
                try:
                    parentdir, name = self.get_parentdir(arg)
                    inode = parentdir.open(name)
                    if (options.recursive and inode.get_type() == api.S5_TYPE_DIR):
                        self._dir_clear(inode)
                        parentdir.rmdir(name)
                    else:
                        parentdir.unlink(name)
                except api.S5fsException as e:
                    self._parse_rm.error(str(e))
             
    def help_rm(self):
        self._parse_rm.print_help()
       
    def complete_rm(self, text, line, begin, end):
        return self.filepath_completion(text, line, begin, end)

    def do_rmdir(self, args):
        try:
            (options, args) = self._parse_rmdir.parse_args(shlex.split(args))
        except ValueError as e:
            self._parse_rmdir.error(str(e))
            return

        if (len(args) == 0):
            self._parse_rmdir.error("command requires at least one argument")
        else:
            for arg in args:
                try:
                    parentdir, name = self.get_parentdir(arg)
                    parentdir.rmdir(name)
                except api.S5fsException as e:
                    self._parse_rmdir.error(str(e))
                
    def help_rmdir(self):
        self._parse_rmdir.print_help()

    def complete_rmdir(self, text, line, begin, end):
        return self.filepath_completion(text, line, begin, end, types=set([ api.S5_TYPE_DIR ]))

    def do_touch(self, args):
        try:
            (options, args) = self._parse_touch.parse_args(shlex.split(args))
        except ValueError as e:
            self._parse_touch.error(str(e))
            return

        if (len(args) == 0):
            self._parse_touch.error("command requires at least one argument")
        else:
            for arg in args:
                try:
                    parentdir, name = self.get_parentdir(arg)
                    if (parentdir.open(name) == None):
                        parentdir.create(name)
                except api.S5fsException as e:
                    self._parse_touch.error(str(e))

    def help_touch(self):
        self._parse_touch.print_help()

    def complete_touch(self, text, line, begin, end):
        return self.filepath_completion(text, line, begin, end)

    def do_mkdir(self, args):
        try:
            (options, args) = self._parse_mkdir.parse_args(shlex.split(args))
        except ValueError as e:
            self._parse_mkdir.error(str(e))
            return

        if (len(args) == 0):
            self._parse_mkdir.error("command requires at least one argument")
        else:
            for arg in args:
                try:
                    parentdir, name = self.get_parentdir(arg)
                    parentdir.mkdir(name)
                except api.S5fsException as e:
                    self._parse_mkdir.error(str(e))

    def help_mkdir(self):
        self._parse_mkdir.print_help()

    def complete_mkdir(self, text, line, begin, end):
        return self.filepath_completion(text, line, begin, end)

    def getfile(self, source, dest):
        dest.truncate()
        
        loc = 0
        data = source.read(20000)
        while (len(data) != 0):
            dest.write(loc, data)
            loc += len(data)
            data = source.read(20000)
        source.close()

    def do_getfile(self, args):
        try:
            (options, args) = self._parse_getfile.parse_args(shlex.split(args))
        except ValueError as e:
            self._parse_getfile.error(str(e))
            return

        if (len(args) != 2):
            self._parse_getfile.error("command requires source and destination paths")
        else:
            try:
                source = open(args[0], 'r')
                dest = self.open(args[1], create=True)
                self.getfile(source, dest)
            except api.S5fsException as e:
                self._parse_getfile.error(str(e))
            except IOError as e:
                self._parse_getfile.error(str(e))

    def help_getfile(self):
        self._parse_getfile.print_help()

    def complete_getfile(self, text, line, begin, end):
        argnum = self.completion_argnum(text, line, begin, end)
        if (argnum == 1):
            try:
                return self.real_filepath_completion(text, line, begin, end)
            except Exception as e:
                print str(e)
        elif (argnum == 2):
            return self.filepath_completion(text, line, begin, end, types=set([ api.S5_TYPE_DATA ]))
        else:
            return []

    def do_putfile(self, args):
        try:
            (options, args) = self._parse_putfile.parse_args(shlex.split(args))
        except ValueError as e:
            self._parse_putfile.error(str(e))
            return

        if (len(args) != 2):
            self._parse_putfile.error("command requires a source and destination")
        else:
            try:
                source = self.open(args[0])
                if (None == source):
                    self._parse_putfile.error("no such file: {0}".format(args[0]))
                dest = open(args[1], 'w+')
                
                loc = 0
                data = source.read(loc, 20000)
                while (len(data) != 0):
                    dest.write(data)
                    loc += len(data)
                    data = source.read(loc, 20000)
                dest.close()
            except api.S5fsException as e:
                self._parse_getfile.error(str(e))
            except IOError as e:
                self._parse_getfile.error(str(e))

    def help_putfile(self):
        self._parse_putfile.print_help()

    def complete_putfile(self, text, line, begin, end):
        argnum = self.completion_argnum(text, line, begin, end)
        if (argnum == 1):
            return self.filepath_completion(text, line, begin, end, types=set([ api.S5_TYPE_DATA ]))
        elif (argnum == 2):
            return self.real_filepath_completion(text, line, begin, end)
        else:
            return []

    def do_format(self, args):
        try:
            (options, args) = self._parse_format.parse_args(shlex.split(args))
        except ValueError as e:
            self._parse_format.error(str(e))
            return

        if (options.size == None and options.blocks == None):
            self._parse_format.error("must specify either -s or -b option to give size of formatted disk")
        elif (options.size != None and options.blocks != None and options.size / api.S5_BLOCK_SIZE != options.blocks):
            self._parse_format.error("specified conflicting -s and -b options, please only specify one")
        elif (options.inodes == None):
            self._parse_format.error("must specify -i option to give number of inodes to reserve on disk")
        else:
            if (options.size != None):
                size = options.size
            else:
                size = options.blocks * api.S5_BLOCK_SIZE
            self._simdisk.format(options.inodes, size)

        if (options.directory):
            q = Queue.Queue()
            q.put(".")
            while(not q.empty()):
                curr = q.get()
                real = os.path.join(options.directory, curr)
                mode = os.stat(real).st_mode
                if (stat.S_ISDIR(mode)):
                    for path in os.listdir(real):
                        q.put(os.path.join(curr, path))
                    parent, name = self.get_parentdir(os.path.join("/", curr))
                    if ("." != name):
                        parent.mkdir(name)
                else:
                    source = open(real, 'r')
                    dest = self.open(os.path.join("/", curr), create=True)
                    self.getfile(source, dest)

    def default(self, line):
        if (line.strip() == "EOF"):
            print("\n")
            return True
        else:
            print(line.split()[0] + " is not a valid command")
            return False

_parser = optparse.OptionParser(usage="usage: %prog <simdisk> [options]", add_help_option=False,
                                description="command line tool to manipulate S5FS simdisks for Weenix")
_parser.add_option("-e", "--execute", action="append", default=[], metavar="COMMAND",
                   help="Executes the given command on the simdisk. This flag can be used multiple times, the commands will be executed "
                   "in the order they appear on the command line. To see a list of commands use the -h flag. For information on a "
                   "specific command use the -c <command> flag.")
_parser.add_option("-h", "--help", action="store_true", default=False,
                   help="Prints help information and a list of fsmaker commands.")
_parser.add_option("-c", "--command", action="append", default=[],
                   help="Prints usage information about the specified command.")
_parser.add_option("-i", "--interactive", action="store_true", default=False,
                   help="Causes fsmaker to open an interactive shell after executing all commands specified with the -e option.")
(options, args) = _parser.parse_args()

try:
    if (len(args) > 1):
        _parser.error("command takes at most one positional argument, but {0} were given".format(len(args)))
    if (len(args) < 1 and not options.help and len(options.command) == 0):
        _parser.error("command requires a simdisk file path")

    if (len(args) < 1):
        fs = FsmakerShell(api.Simdisk(tempfile.TemporaryFile()))
    else:
        try:
            fs = FsmakerShell(api.Simdisk(open(args[0], 'rb+')))
        except IOError as e:
            if (e.errno == errno.ENOENT):
                fs = FsmakerShell(api.Simdisk(open(args[0], 'wb+')))
            else:
                raise e

    if (options.help):
        _parser.print_help()
        fs.onecmd("help")
    for command in options.command:
        fs.onecmd("help {0}".format(command))
    for command in options.execute:
        fs.onecmd(command)
    if (options.interactive):
        fs.cmdloop()
except KeyboardInterrupt:
    print("\n")