Skip to content
tools.py 11.9 KiB
Newer Older
Seblu's avatar
Seblu committed
# -*- python -*-
Seblu's avatar
Seblu committed
# -*- coding: utf-8 -*-
# Started 26/05/2011 by Seblu <seblu@seblu.net>

'''
InstallSystems Generic Tools Library
'''

import os
import re
import hashlib
Seblu's avatar
Seblu committed
import shutil
import urllib2
import time
Sebastien Luttringer's avatar
Sebastien Luttringer committed
from subprocess import call, check_call, CalledProcessError

from installsystems.progressbar import ProgressBar, Percentage, FileTransferSpeed
from installsystems.progressbar import Bar, BouncingBar, ETA, UnknownLength
from installsystems.tarball import Tarball
from installsystems.printer import *

Seblu's avatar
Seblu committed

################################################################################
# Classes
################################################################################

class PipeFile(object):
Seblu's avatar
Seblu committed
    '''
    Pipe file object if a file object with extended capabilties
    like printing progress bar or compute file size, md5 on the fly
Seblu's avatar
Seblu committed
    '''

    def __init__(self, path=None, mode="r", fileobj=None, timeout=3,
                 progressbar=False):
        self.open(path, mode, fileobj, timeout)
        # start progressbar display if asked
        self.progressbar = progressbar

    def open(self, path=None, mode="r", fileobj=None, timeout=3):
        if path is None and fileobj is None:
            raise AttributeError("You must have a path or a fileobj to open")
        if mode not in ("r", "w"):
            raise AttributeError("Invalid open mode. Must be r or w")
        self.mode = mode
        self.timeout = timeout
        self._md5 = hashlib.md5()
        self.size = 0
        self.mtime = None
        self.consumed_size = 0
        # we already have and fo, nothing to open
        if fileobj is not None:
            self.fo = fileobj
            # seek to 0 and compute filesize if we have and fd
            if hasattr(self.fo, "fileno"):
                self.seek(0)
                self.size = os.fstat(self.fo.fileno()).st_size
        # we need to open the path
        else:
            ftype = pathtype(path)
            if ftype == "file":
                self._open_local(path)
            elif ftype == "http":
                self._open_http(path)
            elif ftype == "ftp":
                self._open_ftp(path)
            elif ftype == "ssh":
                self._open_ssh(path)
                raise IOError("URL type not supported")
        # init progress bar
        # we use 0 because a null file is cannot show a progression during write
        if self.size == 0:
            widget = [ BouncingBar(), " ", FileTransferSpeed() ]
            maxval = UnknownLength
        else:
            widget = [ Percentage(), " ", Bar(), " ", FileTransferSpeed(), " ", ETA() ]
            maxval = self.size
        self._progressbar = ProgressBar(widgets=widget, maxval=maxval)

    def _open_local(self, path):
        '''
        Open file on the local filesystem
        '''
        self.fo = open(path, self.mode)
        sta = os.fstat(self.fo.fileno())
        self.size = sta.st_size
        self.mtime = sta.st_mtime

    def _open_http(self, path):
        '''
        Open a file accross an http server
        '''
        try:
            self.fo = urllib2.urlopen(path, timeout=self.timeout)
        except Exception as e:
            # FIXME: unable to open file
            raise IOError(e)
        # get file size
        if "Content-Length" in self.fo.headers:
            self.size = int(self.fo.headers["Content-Length"])
        else:
            self.size = 0
        # get mtime
        try:
            self.mtime = int(time.mktime(time.strptime(self.fo.headers["Last-Modified"],
                                                       "%a, %d %b %Y %H:%M:%S %Z")))
        except:
            self.mtime = None

    def _open_ftp(self, path):
        '''
        Open file via ftp
        '''
        try:
            self.fo = urllib2.urlopen(path, timeout=self.timeout)
        except Exception as e:
            # FIXME: unable to open file
            raise IOError(e)
        # get file size
        try:
            self.size = int(self.fo.headers["content-length"])
        except:
            self.size = 0

    def _open_ssh(self, path):
        '''
        Open current fo from an ssh connection
        '''
        # try to load paramiko
        try:
            import paramiko
        except ImportError:
            raise IOError("URL type not supported")
        # parse url
        (login, passwd, host, port, path) = re.match(
            "ssh://(([^:]+)(:([^@]+))?@)?([^/:]+)(:(\d+))?(/.*)?", path).group(2, 4, 5, 7, 8)
        if port is None: port = 22
        if path is None: path = "/"
        # open ssh connection
        # we need to keep it inside the object unless it was cutted
        self._ssh = paramiko.SSHClient()
        self._ssh.load_system_host_keys()
        self._ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self._ssh.connect(host, port=port, username=login, password=passwd,
                          look_for_keys=True,
                          timeout=int(self.timeout))
        # swith in sftp mode
        sftp = self._ssh.open_sftp()
        # get the file infos
        sta = sftp.stat(path)
        self.size = sta.st_size
        self.mtime = sta.st_mtime
        # open the file
        self.fo = sftp.open(path, self.mode)
        # this is needed to have correct file transfert speed
        self.fo.set_pipelined(True)

    def close(self):
        if self.progressbar:
            self._progressbar.finish()
        debug("MD5: %s" % self.md5)
        debug("Size: %s" % self.consumed_size)
        self.fo.close()

    def read(self, size=None):
        if self.mode == "w":
            raise IOError("Unable to read in w mode")
        buf = self.fo.read(size)
        length = len(buf)
        self._md5.update(buf)
        self.consumed_size += length
        if self.progressbar and length > 0:
            self._progressbar.update(self.consumed_size)
        return buf

    def flush(self):
        if hasattr(self.fo, "flush"):
            return self.fo.flush()

    def write(self, buf):
        if self.mode == "r":
            raise IOError("Unable to write in r mode")
        self.fo.write(buf)
        length = len(buf)
        self._md5.update(buf)
        self.consumed_size += length
        if self.progressbar and length > 0:
            self._progressbar.update(self.consumed_size)
        return None

    def consume(self, fo=None):
        Consume (read) all data and write it in fo
        if fo is None, data are discarded. This is useful to obtain md5 and size
        Useful to obtain md5 and size
        '''
        if self.mode == "w":
            raise IOError("Unable to read in w mode")
        while True:
            buf = self.read(1048576) # 1MiB
            if len(buf) == 0:
                break
            if fo is not None:
                fo.write(buf)

    @property
    def progressbar(self):
        '''
        Return is progressbar have been started
        '''
        return hasattr(self, "_progressbar_started")

    @progressbar.setter
    def progressbar(self, val):
        '''
        Set this property to true enable progress bar
        '''
        if val == True and not hasattr(self, "_progressbar_started"):
            self._progressbar_started = True
            self._progressbar.start()

    @property
    def md5(self):
        '''
        Return the md5 of read/write of the file
        '''
        return self._md5.hexdigest()

    @property
    def read_size(self):
        '''
        Return the current read size
        '''
        return self.consumed_size

    @property
    def write_size(self):
        '''
        Return the current wrote size
        '''
        return self.consumed_size

################################################################################
# Functions
################################################################################
def smd5sum(buf):
    '''
    Compute md5 of a string
    '''
    m = hashlib.md5()
    m.update(buf)
    return m.hexdigest()

def mkdir(path, uid=None, gid=None, mode=None):
Seblu's avatar
Seblu committed
    '''
    Create a directory and set rights
    '''
Seblu's avatar
Seblu committed
    os.makedirs(path)
    chrights(path, uid, gid, mode)

Seblu's avatar
Seblu committed
def chrights(path, uid=None, gid=None, mode=None, mtime=None):
Seblu's avatar
Seblu committed
    '''
    Set rights on a file
    '''
    if uid is not None:
        os.chown(path, uid, -1)
    if gid is not None:
        os.chown(path, -1, gid)
    if mode is not None:
        os.chmod(path, mode)
Seblu's avatar
Seblu committed
    if mtime is not None:
        os.utime(path, (mtime, mtime))

def pathtype(path):
Seblu's avatar
Seblu committed
    '''
    Return path type. This is usefull to know what kind of path is given
    '''
Seblu's avatar
Seblu committed
    if path.startswith("http://") or path.startswith("https://"):
        return "http"
    if path.startswith("ftp://") or path.startswith("ftps://"):
        return "ftp"
Seblu's avatar
Seblu committed
    elif path.startswith("ssh://"):
        return "ssh"
Seblu's avatar
Seblu committed
    else:
Seblu's avatar
Seblu committed
        return "file"

Seblu's avatar
Seblu committed
def isfile(path):
    '''
    Return True if path is of type file
    '''
    return pathtype(path) == "file"

def abspath(path):
Seblu's avatar
Seblu committed
    '''
    Format a path to be absolute
    '''
    ptype = pathtype(path)
    if ptype in ("http", "ftp", "ssh"):
Seblu's avatar
Seblu committed
        return path
    elif ptype == "file":
        if path.startswith("file://"):
Aurélien Dunand's avatar
Aurélien Dunand committed
            path = path[len("file://"):]
Seblu's avatar
Seblu committed
        return os.path.abspath(path)
    else:
        return None
Seblu's avatar
Seblu committed
def getsize(path):
Seblu's avatar
Seblu committed
    '''
    Get size of a path. Recurse if directory
    '''
Seblu's avatar
Seblu committed
    total_sz = os.path.getsize(path)
    if os.path.isdir(path):
        for root, dirs, files in os.walk(path):
            for filename in dirs + files:
                filepath = os.path.join(root, filename)
                filestat = os.lstat(filepath)
                if stat.S_ISDIR(filestat.st_mode) or stat.S_ISREG(filestat.st_mode):
                    total_sz += filestat.st_size
    return total_sz

def human_size(num):
    '''
    Return human readable size
    '''
    for x in ['Bytes', 'KiB', 'MiB', 'GiB', 'TiB', 'PiB', 'EiB', 'ZiB', 'YiB']:
        if num < 1024.0:
            return "%3.1f%s" % (num, x)
        num /= 1024.0
    return "%3.1f%s" % (num, x)
Sebastien Luttringer's avatar
Sebastien Luttringer committed

def chroot(path, shell="/bin/bash", mount=True):
    '''
    Chroot inside a directory and call shell
    if mount is true, mount /{proc,dev,sys} inside the chroot
    '''
    # try to guest distro
    if os.path.exists(os.path.join(path, "etc/debian_version")):
        distro="debian"
    elif os.path.exists(os.path.join(path, "etc/arch-release")):
        distro="archlinux"
    else:
        distro=None
    # try to mount /proc /sys /dev
    if mount:
        mps = ("proc", "sys", "dev", "dev/pts")
        arrow("Mouting filesystems")
        for mp in mps:
            origin =  "/%s" % mp
            target = os.path.join(path, mp)
            if os.path.ismount(origin) and os.path.isdir(target):
                arrow("%s -> %s" % (origin, target), 1)
                try:
                    check_call(["mount",  "--bind", origin, target], close_fds=True)
                except CalledProcessError as e:
                    warn("Mount failed: %s.\n" % e)
    # in case of debian disable policy
    if distro == "debian":
        arrow("Creating debian chroot housekeepers")
        # create a chroot header
        try: open(os.path.join(path, "etc/debian_chroot"), "w").write("CHROOT")
        except: pass
        # fake policy-d
        try: open(os.path.join(path, "usr/sbin/policy-rc.d"), "w").write("#!/bin/bash\nexit 42\n")
        except: pass
    # chrooting
    arrow("Chrooting inside %s and running %s" % (path, shell))
    call(["chroot", path, shell], close_fds=True)
    # cleaning debian stuff
    if distro == "debian":
        arrow("Removing debian chroot housekeepers")
        for f in ("etc/debian_chroot", "usr/sbin/policy-rc.d"):
            try: os.unlink(os.path.join(path, f))
            except: pass
    # unmounting
    if mount:
        arrow("Unmouting filesystems")
        for mp in reversed(mps):
            target = os.path.join(path, mp)
            if os.path.ismount(target):
                arrow(target, 1)
                call(["umount", target], close_fds=True)