#!/usr/bin/env python

""" save.py

- Partitions are identified by their labels
- The backup partition is labelled 'Saved'
- Each backup is in a directory
    <label>/
        archive.img for partclone backups
        ...         for rsync backups
- A .exclude file at the root of a partition prevents
  rsync from saving some files (see rsync man)
"""

SAVED = 'Saved'
EXCLUDE = '.exclude'

import curses, os, re, subprocess, sys, time

PATH = "/usr/bin:/bin:/usr/sbin:/sbin"
DEPENDENCIES = """
    blkid
    mount umount ntfs-3g
    rsync
"""

def copyright(scr, X, Y):
    msg = "'save' (C) Christophe Delord - www.cdsoft.fr"
    scr.addstr(Y-1, X-len(msg)-1, msg, curses.A_DIM)

class Error(Exception):
    def __init__(self, msg): self.msg = msg
    def __str__(self): return "Error: "+self.msg

class UI:

    def __init__(self, scr):
        self.scr = scr

    def open(self, title):
        self.Y, self.X = self.scr.getmaxyx()
        self.menu = []
        self.attr()
        self.scr.clear()
        self.scr.addstr(0, 0, title.center(self.X), curses.A_REVERSE|curses.A_BOLD)
        self.scr.addstr("\n")
        self.scr.refresh()

    def attr(self, attr=curses.A_NORMAL):
        self._attr = attr

    def write(self, txt):
        self.scr.addstr(txt, self._attr)
        self.scr.refresh()

    def add(self, func, txt):
        self.menu.append((func, txt))

    def select(self):
        y, x = self.scr.getyx()
        width = max(len(txt) for (func,txt) in self.menu)
        selected = 0
        while True:
            self.scr.addstr("\n")
            for i, (func, txt) in enumerate(self.menu):
                if i == selected:
                    attr = curses.A_REVERSE
                    left, right = "==> ", " <=="
                else:
                    attr = curses.A_NORMAL
                    left, right = "    ", "    "
                self.scr.addstr(left)
                self.scr.addstr(" "+txt.ljust(width)+" ", attr)
                self.scr.addstr(right+"\n")
            self.scr.addstr("\n")
            copyright(self.scr, self.X, self.Y)
            ch = self.scr.getch()
            if ch == curses.KEY_DOWN: selected = min(selected+1, len(self.menu)-1)
            if ch == curses.KEY_UP: selected = max(selected-1, 0)
            if ch == 27: return None
            if ch == 10: return self.menu[selected][0]
            self.scr.move(y, x)

class Part:

    def __init__(self, dev, type_, label, size):
        self.dev = dev
        self.type = type_
        self.label = label
        self.size = size
        system(["umount", self.dev], devnull, devnull)
        for mounted in open('/etc/mtab', 'rt'):
            mounted, mountpoint, _ = mounted.split(" ", 2)
            if dev == mounted:
                self.mountpoint = mountpoint
                self.already_mounted = True
                self.mounted = True
                break
        else:
            self.mountpoint = "/tmp/%s"%self.label
            self.already_mounted = False
            self.mounted = False
        print "    %s"%self

    def __str__(self):
        return "%s: %s (%s, %s%s)"%(
            self.dev.ljust(9), self.label.ljust(20), self.type, self.size,
            self.already_mounted and ", ALREADY MOUNTED" or ""
        )

    def is_saved(self):
        return self.label.lower() == SAVED.lower()

    def save(self, save_part):
        print
        print "Save %s (%s) to %s"%(self.label, self.dev, save_part.dev)
        print
        if self.formats == save_part.formats and not self.is_bootable():
            self.save_rsync(save_part)
        else:
            self.save_img(save_part)

    def restore(self, save_part):
        print
        print "Restore %s (%s) from %s"%(self.label, self.dev, save_part.dev)
        print
        save_part.mount(write=False)
        try:
            found = bool(os.listdir(os.path.join(save_part.mountpoint, self.label)))
        except OSError:
            found = False
        if not found:
            raise Error("No backup found for %s (%s)"%(self.dev, self.label))
        save_part.unmount()
        if self.formats == save_part.formats and not self.saved_is_bootable(save_part):
            self.restore_rsync(save_part)
        else:
            self.restore_img(save_part)

    def is_bootable(self):
        bootable = False
        self.mount(write=False)
        try:
            for name in ( 'config.sys', 'CONFIG.SYS', 'autoexec.bat', 'AUTOEXEC.BAT',
                          'Boot', 'bootmgr', 'pagefile.sys', 'hiberfil.sys',
                          'boot', 'dev', 'proc', 'sys' ):
                if os.path.exists(os.path.join(self.mountpoint, name)):
                    bootable = True
        finally:
            self.unmount()
        return bootable

    def saved_is_bootable(self, save_part):
        bootable = False
        save_part.mount(write=False)
        try:
            for name in ('archive.img',):
                if os.path.exists(os.path.join(save_part.mountpoint, self.label, name)):
                    bootable = True
        finally:
            save_part.unmount()
        return bootable

    def mount(self, write=False):
        mount = self.__class__ is NTFS and "ntfs-3g" or "mount"
        rw = write and "rw" or "ro"
        if not self.mounted:
            if not os.path.isdir(self.mountpoint): os.makedirs(self.mountpoint)
            if not system([mount, "-o", rw, self.dev, self.mountpoint]):
                raise Error("Can not mount %s"%self.dev)
            self.mounted = True

    def unmount(self):
        if self.mounted and not self.already_mounted:
            if not system(["umount", self.dev]):
                raise Error("Can not unmount %s"%self.dev)
            self.mounted = False

    def create(self, label):
        assert self.is_saved()
        path = os.path.join(self.mountpoint, label)
        if not os.path.exists(path): os.mkdir(path)
        if not os.path.isdir(path):
            raise Error("Can not create '%s'"%label)

    def save_rsync(self, save_part):
        self.mount(write=False)
        save_part.mount(write=True)
        save_part.create(self.label)
        exclude = os.path.join(self.mountpoint, EXCLUDE)
        opts = ["-aHv", "--delete", "--delete-excluded", "--ignore-errors"]
        if os.path.exists(exclude):
            opts.append("--exclude-from=%s"%exclude)
        if not system(["rsync"]+opts+["%s/"%self.mountpoint, "%s/%s/"%(save_part.mountpoint, self.label)]):
            raise Error("rsync error")

    def restore_rsync(self, save_part):
        self.mount(write=True)
        save_part.mount(write=False)
        opts = ["-aHv", "--delete", "--ignore-errors"]
        if not system(["rsync"]+opts+["%s/%s/"%(save_part.mountpoint, self.label), "%s"%self.mountpoint]):
            raise Error("rsync error")

class Ext(Part):
    formats = ('ext3', 'ext4')

    def save_img(self, save_part):
        self.unmount()
        save_part.mount(write=True)
        save_part.create(self.label)
        if not system("partclone.extfs -c -s %s | pigz > %s/%s/archive.img"%(self.dev, save_part.mountpoint, self.label), shell=True):
            raise Error("partclone error")

    def restore_img(self, save_part):
        self.unmount()
        save_part.mount(write=False)
        if not system("cat %s/%s/archive.img | unpigz | partclone.restore -o %s"%(save_part.mountpoint, self.label, self.dev), shell=True):
            raise Error("partclone error")

class NTFS(Part):
    formats = ('ntfs', )

    def save_img(self, save_part):
        self.unmount()
        save_part.mount(write=True)
        save_part.create(self.label)
        if not system("partclone.ntfs -c -s %s | pigz > %s/%s/archive.img"%(self.dev, save_part.mountpoint, self.label), shell=True):
            raise Error("partclone error")

    def restore_img(self, save_part):
        self.unmount()
        save_part.mount(write=False)
        if not system("cat %s/%s/archive.img | unpigz | partclone.restore -o %s"%(save_part.mountpoint, self.label, self.dev), shell=True):
            raise Error("partclone error")

class VFAT(Part):
    formats = ('vfat', )

    def save_img(self, save_part): self.save_rsync(save_part)

    def restore_img(self, save_part): self.restore_rsync(save_part)

class Partitions:

    types = (Ext, NTFS, VFAT)

    def __init__(self):
        print ">>> Scanning partitions"
        sizes = {}
        for part in open('/proc/partitions'):
            if not part.strip(): continue
            size, dev = part.split()[2:]
            try:
                size = int(size)
            except ValueError:
                pass
            else:
                if size >= 1024*1024: size = "%3d GiB"%(size/1024/1024)
                elif size >= 1024: size = "%3d MiB"%(size/1024)
                else: size = "%3d KiB"%size
                sizes["/dev/"+dev] = size
        self.parts = []
        self.save_part = None
        for block in shell("blkid"):
            dev, attrs = block.split(":", 1)
            attrs = dict(re.findall(r'(\w+)="(.*?)"', attrs))
            if not attrs.get('LABEL'): continue
            if attrs['LABEL'].startswith('.'): continue
            if attrs['LABEL'].startswith('_'): continue
            if attrs['LABEL'].startswith(' '): continue
            for part in Partitions.types:
                if attrs['TYPE'] in part.formats:
                    part = part(dev, attrs['TYPE'], attrs['LABEL'], sizes[dev])
                    if part.is_saved():
                        if self.save_part:
                            raise Error("Two partitions labelled '%s'"%SAVED)
                        self.save_part = part
                    else:
                        if part.label in (p.label for p in self.parts):
                            raise Error("Two partitions labelled '%s'"%part.label)
                        self.parts.append(part)
                    break
        self.parts.sort(key=lambda part: (part.dev[:8], int(part.dev[8:] or 0)))
        if not self.save_part:
            raise Error("No partition labelled '%s'."%SAVED)

    def unmount(self):
        for part in self.parts:
            part.unmount()
        self.save_part.unmount()

    def disp(self):
        ui.attr(curses.A_BOLD)
        print >>ui, "Saveable partitions:"
        ui.attr()
        print >>ui
        for part in self.parts:
            print >>ui, "     %s"%part
        print >>ui
        ui.attr(curses.A_BOLD)
        print >>ui, "Backup partition:"
        ui.attr()
        print >>ui
        print >>ui, "     %s"%self.save_part
        print >>ui

    def menu(self, ui, title):
        print >>ui, " "*((ui.X-len(title))/2-1),
        ui.attr(curses.A_BOLD)
        print >>ui, title.capitalize()
        ui.attr()
        for part in self.parts:
            ui.add(part, str(part))
        ui.add(None, "Quit")

def shell(*cmd):
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE)
    stdout = p.communicate()[0]
    if p.returncode != 0: return []
    return stdout.splitlines()

def system(cmd, stdout=None, stderr=None, shell=None):
    sys.stdout.flush()
    p = subprocess.Popen(cmd, stdout=stdout, stderr=stderr, shell=shell)
    p.communicate()
    return p.returncode == 0
devnull = open('/dev/null')

def check_root():
    if os.getuid() != 0:
        sys.exit(os.system("sudo %s"%sys.argv[0]))

def check_dependencies():
    # PATH entries that are sometimes missing
    os_PATH = os.environ['PATH'].split(':')
    for path in PATH.split(':'):
        if path not in os_PATH:
            os.environ['PATH'] = os.environ['PATH'] + ":" + path
    # check for dependencies in PATH
    missing = 0
    for name in DEPENDENCIES.split():
        for path in os.environ['PATH'].split(':'):
            fullname = os.path.join(path, name)
            if os.path.isfile(fullname):
                break
        else:
            print "%s not found"%name
            missing += 1
    if missing > 0:
        sys.exit(-1)

def select(title):
    ui.open(title.capitalize())
    parts.menu(ui, "Partition to %s:"%title)
    part = ui.select()
    return part

def Save():
    part = select('save')
    if part: return lambda: part.save(parts.save_part)

def Restore():
    part = select('restore')
    if part: return lambda: part.restore(parts.save_part)

def shutdown(cmd):
    os.system(cmd)
    while True: time.sleep(60)

def Reboot():
    return lambda: shutdown('shutdown -r now')

def Halt():
    return lambda: shutdown('shutdown -h now')

def main_menu(scr):
    global ui
    ui = UI(scr)
    ui.open("Save / Restore")
    parts.disp()
    ui.add(Save,    "Save a partition")
    ui.add(Restore, "Restore a partition")
    if "".join(shell("uname", "-n")).strip() == "sysresccd":
        ui.add(Reboot, "Reboot")
        ui.add(Halt,   "Halt")
    else:
        ui.add(None,   "Quit")
    func = ui.select()
    return func and func()

if __name__ == '__main__':

    check_root()
    check_dependencies()

    try:
        parts = None
        parts = Partitions()
        func = curses.wrapper(main_menu)
        if func: func()

    except Error, e:
        print >>sys.stderr, e

    finally:
        if parts: parts.unmount()