path: root/forkcopy-server.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sys, os
import curses
import socket
import threading
import time
import shutil
import tempfile
import getpass

scr = None
logfile = None

CHUNK_SIZE = 1048576

def log(s):
    print >> logfile, s
    logfile.flush()

def logerr(s):
    log('ERROR: ' + s)

class MainLoopInterface:
    def main_loop(self): raise NotImplementedError()
    def should_quit(self): raise NotImplementedError()

class RemoteControlInterface:
    def next(self): raise NotImplementedError()

def convert_bytes(bytes):
    bytes = float(bytes)
    if bytes >= 1099511627776:
        terabytes = bytes / 1099511627776
        size = '%.2fT' % terabytes
    elif bytes >= 1073741824:
        gigabytes = bytes / 1073741824
        size = '%.2fG' % gigabytes
    elif bytes >= 1048576:
        megabytes = bytes / 1048576
        size = '%.2fM' % megabytes
    elif bytes >= 1024:
        kilobytes = bytes / 1024
        size = '%.2fK' % kilobytes
    else:
        size = '%.2fb' % bytes
    return size

class RemoteControl(RemoteControlInterface, threading.Thread):

    def __init__(self):
        super(RemoteControl, self).__init__()
        self._lines = []
        self._lines_lock = threading.Lock()
        self.daemon = True

    def next(self):
        if not self._lines_lock.acquire(False):
            return None
        else:
            line = None
            try:
                if len(self._lines):
                    line = self._lines.pop()
            finally:
                self._lines_lock.release()
            return line

    def run(self):
        sock_path = os.path.join(tempfile.gettempdir(),
                'forkcopy-socket_' + getpass.getuser())
        sock_server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        if os.path.exists(sock_path):
            os.remove(sock_path)
        sock_server.bind(sock_path)
        sock_server.listen(1)
        while True:
            sock, addr = sock_server.accept()
            data = sock.recv(4096)
            buf = ''
            while len(data) != 0:
                buf += data
                data = sock.recv(4096)
            sock.send('OK')
            sock.close()
            with self._lines_lock:
                self._lines.extend(buf.splitlines())
        sock_server.close()
        logerr('should not be here deverror')

class forkcopy:

    ## @param main_loop a MainLoopInterface
    def __init__(self, main_loop):
        self._chunk_size = CHUNK_SIZE
        self._size_copied = 0
        self._is_active = False
        self._main_loop = main_loop

    # note I don't not care about free space, code here will raise IOError if
    # an error arise.
    def copy(self, ifn, ofn):
        if os.path.exists(ofn) or os.path.exists(ofn + '.part'):
            logerr('file `%s\' or `%s.part\' already exists, skipping' %
                    (ofn, ofn))
            return
        self._is_active = True
        self._size_copied = 0
        try:
            dst = os.path.dirname(ofn)
            if not os.path.exists(dst):
                os.makedirs(dst)
            remove_ofn = False
            with open(ofn + '.part', 'wb') as target, \
                    open(ifn, 'rb') as source:
                data = source.read(self._chunk_size)
                while len(data) != 0:
                    target.write(data)
                    self._size_copied += len(data)
                    self._main_loop.main_loop()
                    if self._main_loop.should_quit():
                        remove_ofn = True
                        break
                    data = source.read(self._chunk_size)
            if remove_ofn:
                os.remove(ofn + '.part')
            else:
                os.rename(ofn + '.part', ofn)
            try:
                shutil.copystat(ifn, ofn)
            #except OSError, why:
            except OSError:
                # can't copy file access times on Windows and windows fs
                #logerr('can\'t copy attributes, %s' % str(why))
                pass
        finally:
            self._is_active = False

    def is_active(self): return self._is_active
    def copied_size(self) : return self._size_copied
    def reset_size(self): self._size_copied = 0



class MainLoop(MainLoopInterface):

    def __init__(self):
        self._rc   = None
        self._fc   = None
        self._quit = False

        self._currentbase    = '.'
        self._filelist       = []
        self._currentfile    = (None, 0, None, None)
        self._previousfile   = False
        self._previousstatus = False

        self._currentfilecopiedsize = 0
        self._copiedsize = 0
        self._previouscopiedsize = 0
        self._totalsize = 0

        self._speed = 0
        self._speed_timer = None
        self._speed_timertime = 5
        self._speed_lastsize = 0
        self._speed_lock = threading.Lock()

        self._fc = forkcopy(self)
        self._rc = RemoteControl()
        self._rc.start()
        self._speed_timer = threading.Timer(self._speed_timertime,
                self._determine_speed)
        self._speed_timer.start()

        curses.curs_set(0)
        scr.nodelay(1)

        scr.addstr(
                'Filename :\n' +
                'File size (Copied/Total):\n' +
                'Current destination:\n' +
                'Base destination directory:\n' +
                'Copy size (Copied/Total):\n' +
                'Current speed:\n' +
                'Current progress:\n' +
                'Global progress:\n')

        scr.refresh()

    def should_quit(self): return self._quit

    def _determine_speed(self):
        self._speed_timer = threading.Timer(self._speed_timertime,
                self._determine_speed)
        self._speed_timer.start()
        with self._speed_lock:
            self._speed = (self._copiedsize - self._speed_lastsize) / \
                        self._speed_timertime
            self._speed_lastsize = self._copiedsize

    def _manage_copy(self):
        if self._fc.is_active():
            return
        # add new files if available in the list
        if len(self._filelist) != 0:
            self._currentfile = self._filelist.pop()
            ifn, size, ofn, base = self._currentfile
            fcsize = self._fc.copied_size()
            self._fc.reset_size()
            self._copiedsize += fcsize
            self._fc.copy(ifn, ofn)
        else:
            self._currentfile = ('Waiting for new files to arrive',
                    0, '', self._currentbase)
            time.sleep(1) # avoid wasting cpu cycles when waiting for files

    def _print_progressbar(self, percent, y):
        if percent < 0 or percent > 100: percent = 0
        height, width = scr.getmaxyx()
        width -= 2 # substract '[]' characters
        width += 1 # take whole line (bad sums below...)
        if width <= 4: return
        scr.addstr(y, 0, '[')
        nbsharp = width*percent/100
        for i in range(1, nbsharp):
            scr.addstr(y, i, '#')
        for i in range(1 + nbsharp, width):
            scr.addstr(y, i, ' ')
        scr.addstr(']\n')

    def _normalize_str(self, s, size):
        if len(s) <= 5 or size <= 5:
            return s
        if len(s) >= size:
            return '...' + s[len(s)-size-3:]
        return s

    def _manage_display(self):
        screenh, screenw = scr.getmaxyx()
        fcsize = self._fc.copied_size()
        self._fc.reset_size()
        self._currentfilecopiedsize += fcsize
        self._copiedsize += fcsize
        if (self._currentfile != self._previousfile or
                self._copiedsize != self._previouscopiedsize):
            if self._currentfile != self._previousfile:
                self._currentfilecopiedsize = fcsize
            self._previousfile = self._currentfile
            self._previouscopiedsize = self._copiedsize
            (ifn, size, ofn, base) = self._currentfile
            speed = 0
            with self._speed_lock:
                speed = self._speed
            perc = 0
            gperc = 0
            if size != 0:
                perc = self._currentfilecopiedsize * 100 / size
            if self._totalsize != 0:
                gperc = self._copiedsize * 100 / self._totalsize
            scr.addstr(0, 28, self._normalize_str(ifn, screenw - 28) + '\n')
            scr.addstr(1, 28, convert_bytes(self._currentfilecopiedsize)
                    + '/' + convert_bytes(size) + '\n')
            scr.addstr(2, 28, self._normalize_str(ofn, screenw - 28) + '\n')
            scr.addstr(3, 28, self._normalize_str(base, screenw - 28) + '\n')
            scr.addstr(4, 28, convert_bytes(self._copiedsize) + '/' +
                    convert_bytes(self._totalsize) + '\n')
            scr.addstr(5, 28, convert_bytes(speed) + '/s\n')
            scr.addstr(6, 28, str(perc) + '%\n')
            scr.addstr(7, 28, str(gperc) + '%\n')
            self._print_progressbar(perc, 8)
            self._print_progressbar(gperc, 9)
            scr.refresh()

    def _gettreefiles(self, src, prefix = None):
        if src[len(src)-1] == os.sep:
            src = src[0:len(src)-1] # avoid basename troubles
        names = os.listdir(src)
        dstbase = ''
        if prefix is None:
            dstbase = os.path.join(self._currentbase, os.path.basename(src))
        else:
            dstbase = os.path.join(prefix, os.path.basename(src))
        ret = []
        for name in names:
            srcname = os.path.join(src, name)
            dstname = os.path.join(dstbase, name)
            try:
                if os.path.isdir(srcname):
                    ret.extend(self._gettreefiles(srcname, dstname))
                else:
                    size = 0
                    try:
                        size = os.path.getsize(srcname)
                    except os.error:
                        logerr('`%s\' inaccessible or command error.' %
                                srcname)
                        continue
                    newfile = (srcname, size, dstname, self._currentbase)
                    ret.append(newfile)
                    self._totalsize += size
            except (IOError, os.error), why:
                logerr('`%s\', %s.' % (srcname, str(why)))
            except shutil.Error, err:
                logerr('`%s\', %s.' % (srcname, str(err.args)))

        return ret

    def _command_add(self, ifn):
        if not os.path.isdir(ifn):
            size = 0
            ofn = os.path.join(self._currentbase,
                    os.path.basename(ifn))
            try:
                size = os.path.getsize(ifn)
            except os.error:
                logerr('`%s\' inaccessible or command error.' %
                        ifn)
                return
            newfile = (ifn, size, ofn, self._currentbase)
            log('adding file: ' + str(newfile))
            self._filelist.append(newfile)
            self._totalsize += size
        else:
            newfiles = self._gettreefiles(ifn)
            self._filelist.extend(newfiles)
            log('adding files: ' + str(newfiles))

    def _command_iadd(self, arg):
        logerr('iadd not implemented')

    def _declare_quit(self):
        self._quit = True
        if self._speed_timer is not None:
            self._speed_timer.cancel()
            self._speed_timer = None

    def main_loop(self):
        if self._quit:
            return # app is quitting, do not do any operation

        item = self._rc.next()
        if item != None:
            index = item.find(' ')
            command = ''
            arg = ''
            if index != -1:
                command = item[0:index]
                arg = item[index + 1:]
            else:
                command = item
            log('command: `%s\'\narg: `%s\'' % (command, arg))
            if command == 'add':
                self._command_add(arg)
            elif command == 'iadd': # add as immutable, copy dir structure
                self._command_iadd(arg)
            elif command == 'quit':
                self._declare_quit()
                return
            else:
                logerr('unrecognized command: `%s\' on line `%s\'' %
                        (command, item))

        if not self._rc.is_alive():
            self._declare_quit()
            return

        key = scr.getch()
        if key != curses.ERR:
            if key == ord('q') or key == ord('Q'):
                self._declare_quit()
            elif key == ord('p'):
                scr.addstr(23, 0, 'Paused : touch the \'p\' key to continue')
                scr.refresh()
                while True:
                    key = scr.getch()
                    if key != curses.ERR and key == ord('p'):
                        break
                    else:
                        time.sleep(1)
                scr.addstr(23, 0, '                                         ')
                scr.refresh()

        self._manage_copy()
        self._manage_display()

def main(stdscr):
    global scr
    scr = stdscr
    main_loop = MainLoop()
    while not main_loop.should_quit():
        main_loop.main_loop()

if __name__ == '__main__':
    # redirecting stderr errors to stderr.txt since it's hard to see inside
    # alternative curses screen.
    stderr = open(os.path.join(tempfile.gettempdir(),
        'forkcopy-stderr_' + getpass.getuser() + '.txt'), 'wb')
    sys.stderr.close()
    sys.stderr = stderr
    with open(os.path.join(tempfile.gettempdir(),
            'forkcopy-log_' + getpass.getuser() + '.txt'), 'wb') as logfile:
        log('START')
        curses.wrapper(main)
        log('STOP')
    sys.exit(0)