Skip to content

Instantly share code, notes, and snippets.

@fuzzy

fuzzy/dd.py

Created May 3, 2015 22:52
Show Gist options
  • Save fuzzy/f0e47cc3f45014cee3a7 to your computer and use it in GitHub Desktop.
Save fuzzy/f0e47cc3f45014cee3a7 to your computer and use it in GitHub Desktop.
Download ZIP
dd.py
Raw
dd.py
#!/usr/bin/env python
###########
# Imports #
###########
# Stdlib
import re
import os
import sys
import time
import urllib2
import Queue
import threading
#####################
# Global constructs #
#####################
sigQueue = Queue.Queue() # Signal queue
logQueue = Queue.Queue() # Log message queue
##################
# Output Helpers #
##################
def terminal_size():
import fcntl, termios, struct
h, w, hp, wp = struct.unpack('HHHH',
fcntl.ioctl(0, termios.TIOCGWINSZ,
struct.pack('HHHH', 0, 0, 0, 0)))
return w, h
def humanSize(bytes=0):
kbyte = 1024
mbyte = (kbyte**2)
gbyte = (kbyte**3)
tbyte = (kbyte**4)
pbyte = (kbyte**5)
ebyte = (kbyte**6)
zbyte = (kbyte**7)
if bytes < kbyte:
retv = '%dB' % int(bytes)
return unicode('%9s' % retv)
elif bytes >= kbyte and bytes < mbyte:
retv = '%04.02fKB' % (float(bytes) / float(kbyte))
return unicode('%9s' % retv)
elif bytes >= mbyte and bytes < gbyte:
retv = '%04.02fMB' % (float(bytes) / float(mbyte))
return unicode('%9s' % retv)
elif bytes >= gbyte and bytes < tbyte:
retv = '%04.02fGB' % (float(bytes) / float(gbyte))
return unicode('%9s' % retv)
elif bytes >= tbyte and bytes < pbyte:
retv = '%04.02fTB' % (float(bytes) / float(tbyte))
return unicode('%9s' % retv)
elif bytes >= pbyte and bytes < ebyte:
retv = '%04.02fPB' % (float(bytes) / float(pbyte))
return unicode('%9s' % retv)
elif bytes >= ebyte and bytes < zbyte:
retv = '%04.02fEB' % (float(bytes) / float(ebyte))
return unicode('%9s' % retv)
else:
retv = '%04.02fZB' % (float(bytes) / float(zbyte))
return unicode('%9s' % retv)
def humanTime(seconds=0):
# These are for convenience
minute = 60
hour = (minute**2)
day = (hour*24)
week = (day*7)
month = (week*4)
year = (month*12)
secs, mins, hrs, days, weeks, months, years = 0, 0, 0, 0, 0, 0, 0
if seconds > year:
years = (seconds / year)
tmp = (seconds % year)
seconds = tmp
if seconds > month:
months = (seconds / month)
tmp = (seconds % month)
seconds = tmp
if seconds > week:
weeks = (seconds / week)
tmp = (seconds % week)
seconds = tmp
if seconds > day:
days = (seconds / day)
tmp = (seconds % day)
seconds = tmp
if seconds > hour:
hrs = (seconds / hour)
tmp = (seconds % hour)
seconds = tmp
if seconds > minute:
mins = (seconds / minute)
secs = (seconds % minute)
if seconds < minute:
secs = seconds
if years != 0:
return unicode('%4dy%2dm%1dw%1dd %02d:%02d:%02d' % (
years, months, weeks, days, hrs, mins, secs
))
if months != 0:
return unicode('%2dm%1dw%1dd %02d:%02d:%02d' % (
months, weeks, days, hrs, mins, secs
))
if weeks != 0:
return unicode('%1dw%1dd %02d:%02d:%02d' % (
weeks, days, hrs, mins, secs
))
if days != 0:
return unicode('%1dd %02d:%02d:%02d' % (days, hrs, mins, secs))
return unicode('%02d:%02d:%02d' % (hrs, mins, secs))
class Scale(str):
def __len__(self):
'''
This ensures that you will always get the actual length of the string,
minus the extended characters. Which is of course important, when you
are calculating output field sizes. Requires the re module.
'''
tmp = self[:]
cnt = 0
for i in re.sub('\\x1b[\[0-9;]*m', '', tmp):
cnt += 1
return(cnt)
def __getattr__(self, method):
'''
This is essentially an implimentation of Ruby's .method_missing
that shortens the code dramatically, and allows for simply extending
to support other escape codes. As a note, the modifier methods like
.bold() and .underline() and such, need to come before the color
methods. The color should always be the last modifier.
'''
method_map = {
'black': {'color': True, 'value': 30, 'mode': 'm'},
'red': {'color': True, 'value': 31, 'mode': 'm'},
'green': {'color': True, 'value': 32, 'mode': 'm'},
'yellow': {'color': True, 'value': 33, 'mode': 'm'},
'blue': {'color': True, 'value': 34, 'mode': 'm'},
'purple': {'color': True, 'value': 35, 'mode': 'm'},
'cyan': {'color': True, 'value': 36, 'mode': 'm'},
'white': {'color': True, 'value': 37, 'mode': 'm'},
'clean': {'color': False, 'value': 0, 'mode': 'm'},
'bold': {'color': False, 'value': 1, 'mode': 'm'},
'underline': {'color': False, 'value': 4, 'mode': 'm'},
'blink': {'color': False, 'value': 5, 'mode': 'm'},
'reverse': {'color': False, 'value': 7, 'mode': 'm'},
'conceal': {'color': False, 'value': 8, 'mode': 'm'},
}
def get(self, **kwargs):
if method_map[method]['color']:
reset=''
else:
reset=''
return(
Scale('%s[%s%s%s%s' % (
reset,
method_map[method]['value'],
method_map[method]['mode'],
self,
reset
)
))
if method in method_map:
return get.__get__(self)
else:
raise(AttributeError, method)
class Logger(threading.Thread):
def __init__(self, debug=True, quiet=False):
# Inititalize ourselves as a thread
threading.Thread.__init__(self)
# Setup our signal and log queues
global sigQueue
global logQueue
self.sigQueue = sigQueue
self.logQueue = logQueue
# set our display flags
self.debug = debug
self.quiet = quiet
# define our tags
self.tags = {'info': Scale('INFO').bold().green(),
'warn': Scale('WARN').bold().yellow(),
'error': Scale('ERROR').bold().red(),
'fatal': Scale('FATAL').bold().red(),
'debug': Scale('DEBUG').bold().cyan()}
def run(self):
while True:
message = True
while message:
# Do we have any log messages to print?
try:
# Get a message if there is one
message = self.logQueue.get(True, 0.1)
# Lets validate our entry
for t in ('tag', 'data'):
if t not in message.keys():
raise ValueError, 'Improperly formatted log entry: %s' % message.__repr__()
# now lets print our message:
if not self.quiet:
# set a flag to suppress/allow debugging output
printit = True
if message['tag'].lower() == 'debug' and not self.debug:
printit = False
if message['tag'].lower() in self.tags.keys() and printit:
sys.stdout.write('%5s%s %s\n' % (self.tags[message['tag']], Scale(':').bold().white(), message['data']))
sys.stdout.flush()
# Finally, if we have a fatal tag, lets send off our signal, and break outta here
if message['tag'].lower() == 'fatal':
message = False
exitSignal()
except Queue.Empty:
message = False
# And finally if we have recieved a signal, lets exit
try:
a = self.sigQueue.get(True, 0.1)
self.sigQueue.put(True)
return
except Queue.Empty:
pass
def exitSignal(): sigQueue.put(True)
def info(msg): logQueue.put({'tag': 'info', 'data': msg})
def warn(msg): logQueue.put({'tag': 'warn', 'data': msg})
def debug(msg): logQueue.put({'tag': 'debug', 'data': msg})
def error(msg): logQueue.put({'tag': 'error', 'data': msg})
def fatal(msg): logQueue.put({'tag': 'fatal', 'data': msg})
###############
# I/O Objects #
###############
class IOobject:
def __init__(self, uri, mode='r', bsize=1048576):
self.buffsize = bsize
# Now figure out wtf they handed us in the uri
if uri:
if re.match('^(htt|ft).+://.*$', uri):
if uri.find('http') != -1 or uri.find('ftp') != -1:
self.__fp = urllib2.urlopen(uri)
self.size = int(self.__fp.headers['content-length'])
debug('(ht|f)tp url: %s' % self.__fp.__repr__())
elif os.path.isfile(uri) or os.path.islink(uri):
self.__fp = open(uri, mode)
self.size = os.stat(uri).st_size
debug('File: %s' % self.__fp.__repr__())
elif not os.path.isfile(uri) and os.path.isdir('/'.join(uri.split('/')[:-1])):
self.__fp = open(uri, 'w+')
debug('File: %s' % self.__fp.__repr__())
else:
raise ValueError, "Unsupported I/O type. %s" % uri
def read(self):
return self.__fp.read(self.buffsize)
def write(self, data):
return self.__fp.write(data)
class Transfer(threading.Thread):
def __init__(self, ipObj, opObj):
# Initialize the thread stack
threading.Thread.__init__(self)
# setup the signal queue
global sigQueue
self.sigQueue = sigQueue
# And now ourselves
self.ipObj = ipObj
self.opObj = opObj
self.__lock = threading.Lock()
self.__buff = True
self.readBytes = 0
self.wroteBytes = 0
def __wrapOp(self, func):
if func:
self.__lock.acquire()
try: func()
finally: self.__lock.release()
def __recordReadSize(self):
self.readBytes += len(self.__buff)
def __recordWriteSize(self):
self.wroteBytes += len(self.__buff)
def __read(self):
self.__buff = self.ipObj.read()
def __write(self):
self.opObj.write(self.__buff)
def run(self):
debug(str(self.__buff))
while self.__buff:
self.__wrapOp(self.__read)
self.__wrapOp(self.__recordReadSize)
self.__wrapOp(self.__write)
self.__wrapOp(self.__recordWriteSize)
# And finally if we have recieved a signal, lets exit
try:
a = self.sigQueue.get(True, 0.1)
self.sigQueue.put(True)
return
except Queue.Empty:
pass
###################
# Progress output #
###################
def progress(tobj):
start = time.time()
total = tobj.ipObj.size
width = terminal_size()[0]
tobj.start()
while tobj.isAlive():
prc = (float(tobj.wroteBytes) / float(total))
prog = ((float(tobj.wroteBytes) / float(total)) * 100.00)
speed = (float(tobj.wroteBytes) / float(time.time() - start))
# 33 chars out of width are unavailable for hashmarks
hashes = '#'*int((width-33)*prc)
space = '-'*((width - 33) - len(hashes))
sys.stdout.write('\r%s %s%s%s%s %3d%% @ %s/sec' % (
humanSize(tobj.wroteBytes),
Scale('[').bold().blue(),
Scale(hashes).bold().white(),
Scale(space).cyan(),
Scale(']').bold().blue(),
int(prog),
humanSize(speed)))
sys.stdout.flush()
time.sleep(1)
print('')
########
# Main #
########
if __name__ == '__main__':
log = Logger(debug=False)
log.start()
try:
iObj = IOobject('http://ftp.osuosl.org/pub/archlinux/iso/latest/archlinux-2015.05.01-dual.iso')
oObj = IOobject('/', mode='w+')
tObj = Transfer(iObj, oObj)
progress(tObj) #.transfer()
except ValueError, e:
fatal(e)
except:
einfo = sys.exc_info()
print('')
fatal('An unknown error has occurred: %s (line: %d)' % (einfo[1].message, einfo[2].tb_lineno))
debug('exiting')
exitSignal()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment