mdt/files.py

import os
import sys
import time

from mdt import command
from mdt import config
from mdt import console
from mdt import discoverer
from mdt import sshclient


PROGRESS_WIDTH = 45
FILENAME_WIDTH = 30


def MakeProgressFunc(full_filename, width, char='>'):
    def closure(bytes_xferred, total_bytes):
        pcnt = bytes_xferred / total_bytes
        filename = full_filename
        if len(filename) > FILENAME_WIDTH:
            filename = filename[0:FILENAME_WIDTH - 3] + '...'
        left = char * int(pcnt * width)
        right = ' ' * int((1 - pcnt) * width)
        pcnt = '%3d' % (int(pcnt * 100))
        sys.stdout.write('\r{0}% |{1}{2}| {3}'.format(pcnt, left, right, filename))
        sys.stdout.flush()

    return closure


class InstallCommand(command.NetworkCommand):
    def preConnectRun(self, args):
        if len(args) < 2:
            print("Usage: mdt install [<package-filename...>]")
            return False

        return True

    def runWithClient(self, client, args):
        package_to_install = args[1]
        package_filename = os.path.basename(package_to_install)
        remote_filename = os.path.join('/tmp', package_filename)

        sftp_callback = MakeProgressFunc(package_filename, PROGRESS_WIDTH)
        sftp = client.openSftp()
        sftp.put(package_to_install, remote_filename, callback=sftp_callback)
        sftp.close()
        client.close()
        print()

        channel = client.shellExec("sudo /usr/sbin/mdt-install-package {0}; rm -f {0}".format(remote_filename), allocPty=True)
        cons = console.Console(channel, sys.stdin)
        return cons.run()


class PushCommand(command.NetworkCommand):
    def preConnectRun(self, args):
        if len(args) < 3:
            print("Usage: mdt push <filename...> <destination-directory>")
            return False

        for file in args[1:-1]:
            if not os.path.isfile(file):
                print("{0}: Is a directory -- cannot push".format(file))
                return False

        return True

    def runWithClient(self, client, args):
        files_to_push = args[1:-1]
        destination = args[-1]

        try:
            sftp = client.openSftp()
            for file in files_to_push:
                base_filename = os.path.basename(file)
                sftp_callback = MakeProgressFunc(file, PROGRESS_WIDTH)
                remote_filename = os.path.join(destination, base_filename)

                sftp_callback(0, 1)
                sftp.put(file, remote_filename, callback=sftp_callback)
                sftp_callback(1, 1)
                print()
        finally:
            print()
            sftp.close()

        return 0


class PullCommand(command.NetworkCommand):
    def preConnectRun(self, args):
        if len(args) < 3:
            print("Usage: mdt pull [<filename...>]")
            return False

        return True

    def runWithClient(self, client, args):
        files_to_pull = args[1:-1]
        destination = args[-1]

        try:
            sftp = client.openSftp()
            for file in files_to_pull:
                base_filename = os.path.basename(file)
                sftp_callback = MakeProgressFunc(file, PROGRESS_WIDTH, char='<')
                destination_filename = os.path.join(destination, base_filename)

                sftp_callback(0, 1)
                sftp.get(file, destination_filename, callback=sftp_callback)
                sftp_callback(1, 1)
                print()
        finally:
            print()
            sftp.close()

        return 0