Newer
Older
# -*- coding: utf-8 -*-
# Started 26/05/2011 by Seblu <seblu@seblu.net>
'''
InstallSystems Generic Tools Library
'''
import os
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 *
################################################################################
# Classes
################################################################################
class PipeFile(object):
Pipe file object if a file object with extended capabilties
like printing progress bar or compute file size, md5 on the fly
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
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
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)
# 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:
# 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:
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)
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")
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
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:
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):
chrights(path, uid, gid, mode)
def chrights(path, uid=None, gid=None, mode=None, mtime=None):
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)
'''
Return path type. This is usefull to know what kind of path is given
'''
if path.startswith("http://") or path.startswith("https://"):
return "http"
if path.startswith("ftp://") or path.startswith("ftps://"):
return "ftp"
def isfile(path):
'''
Return True if path is of type file
'''
return pathtype(path) == "file"
return path
elif ptype == "file":
if path.startswith("file://"):
return os.path.abspath(path)
else:
return None
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)
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
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)