Newer
Older
# -*- coding: utf-8 -*-
# Started 26/05/2011 by Seblu <seblu@seblu.net>
'''
InstallSystems Generic Tools Library
'''
import os
from progressbar import ProgressBar, Percentage, FileTransferSpeed
from 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
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
def __init__(self, path=None, mode="r", fileobj=None, timeout=3,
progressbar=False):
self.progressbar = progressbar
self.open(path, mode, fileobj, timeout)
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._md5 = hashlib.md5()
self.size = None
self.consumed_size = 0
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.fo = open(path, self.mode)
self.size = os.fstat(self.fo.fileno()).st_size
elif ftype == "http" or ftype == "ftp":
try:
self.fo = urllib2.urlopen(path, timeout=timeout)
except Exception as e:
# FIXME: unable to open file
raise IOError(e)
if "Content-Length" in self.fo.headers:
self.size = int(self.fo.headers["Content-Length"])
else:
raise NotImplementedError
# init progress bar
if self.size is None:
widget = [ BouncingBar(), " ", FileTransferSpeed() ]
maxval = UnknownLength
else:
widget = [ Percentage(), " ", Bar(), " ", FileTransferSpeed(), " ", ETA() ]
maxval = self.size
self._progressbar = ProgressBar(widgets=widget, maxval=maxval)
# start progressbar display if asked
if self.progressbar:
self._progressbar.start()
def close(self):
if self.progressbar:
self._progressbar.finish()
debug("MD5: %s" % self.md5)
debug("Size: %s" % self.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")
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):
'''
Read all data and doesn't save it
Useful to obtain md5 and size
'''
if self.mode == "w":
raise IOError("Unable to read in w mode")
while True:
buf = self.read(65536)
if len(buf) == 0:
break
@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 copy(source, destination, uid=None, gid=None, mode=None, timeout=None):
'''
Copy a source to destination. Take care of path type
'''
stype = pathtype(source)
dtype = pathtype(destination)
# ensure destination is not a directory
if dtype == "file" and os.path.isdir(destination):
destination = os.path.join(destination, os.path.basename(source))
# trivial case
if stype == dtype == "file":
shutil.copy(source, destination)
elif (stype == "http" or stype == "ftp") and dtype == "file":
f_dest = open(destination, "w")
f_source = urllib2.urlopen(source, timeout=timeout)
copyfileobj(f_source, f_dest)
raise NotImplementedError
else:
raise NotImplementedError
# setting destination file rights
if dtype == "file":
chrights(destination, uid, gid, mode)
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)