Merge branch 'wget-requests'

Conflicts:
	pwnlib/util/web.py

pwn/toplevel.py

@@ -26,7 +26,7 @@ from pwnlib.util.misc        import *
 from pwnlib.util.packing     import *
 from pwnlib.util.proc        import pidof
 from pwnlib.util.splash      import *
-from pwnlib.util.web         import wget
+from pwnlib.util.web         import *
 
 # Promote these modules, so that "from pwn import *" will let you access them
 import \

pwnlib/util/misc.py

@@ -44,7 +44,7 @@ def size(n, abbriv = 'B', si = False):
     Convert the length of a bytestream to human readable form.
 
     Args:
-      n(int): The length to convert to human readable form.
+      n(int,str): The length to convert to human readable form
       abbriv(str):
 
     Example:
@@ -59,6 +59,9 @@ def size(n, abbriv = 'B', si = False):
         >>> [size(1024 ** n) for n in range(7)]
         ['1B', '1.00KB', '1.00MB', '1.00GB', '1.00TB', '1.00PB', '1024.00PB']
     """
+    if isinstance(n, str):
+        n = len(n)
+
     base = 1000.0 if si else 1024.0
     if n < base:
         return '%d%s' % (n, abbriv)

pwnlib/util/web.py

@@ -1,18 +1,10 @@
 # -*- coding: utf-8 -*-
-import os, tempfile, urllib2, logging
+import os, tempfile, logging
+from requests import *
+from .misc import size
 log = logging.getLogger(__name__)
 
-# Using powers of ten instead of two,
-# since that's what bandwidth is measured in.
-sizes = (
-    (10**9, 'gB'),
-    (10**6, 'mB'),
-    (10**3, 'kB'),
-    (1,     'B')
-)
-
-
-def wget(url, save=None, timeout=5):
+def wget(url, save=None, timeout=5, **kwargs):
     """wget(url, save=None, timeout=5) -> str
 
     Downloads a file via HTTP/HTTPS.
@@ -33,44 +25,45 @@ def wget(url, save=None, timeout=5):
       >>> result == file('robots.txt').read()
       True
     """
-    response = urllib2.urlopen(url, timeout=timeout);
+    with log.progress("Downloading '%s'" % url) as w:
+        w.status("Making request...")
 
-    if response.code != 200:
-        log.error("Got code %s" % response.code)
-        return
+        response = get(url, stream=True, **kwargs)
 
-    log.waitfor("Downloading '%s'" % url)
-    total_size = response.info().getheader('Content-Length').strip()
-    total_size = int(total_size)
+        if not response.ok:
+            w.failure("Got code %s" % response.status_code)
+            return
 
-    # Find out the next largest size we can represent as
-    for chunk_size, size_name in sizes:
-        if chunk_size < total_size:
-            break
+        total_size = int(response.headers.get('content-length',0))
 
-    # Count chunks as they're received
-    chunks_so_far = 0
-    total_chunks  = total_size / chunk_size
-    total_data    = ''
+        w.status('0 / %s' % size(total_size))
 
-    # Loop until we have all of the data
-    chunk = response.read(chunk_size)
-    while chunk:
-        total_data += chunk
-        chunks_so_far += 1
-        log.status('%s / %s %s' % (chunks_so_far, total_chunks, size_name))
-        chunk = response.read(chunk_size)
+        # Find out the next largest size we can represent as
+        chunk_size = 1
+        while chunk_size < (total_size/10):
+            chunk_size *= 1000
 
-    # Save to the target file if provided
-    if save:
-        if not isinstance(save, (str, unicode)):
-            save = os.path.basename(url)
-            save = save or NamedTemporaryFile(dir='.', delete=False).name
-        with file(save,'wb+') as f:
-            f.write(total_data)
-            log.done_success('Saved data to %r' % f.name)
-    else:
-        log.done_success()
+        # Count chunks as they're received
+        total_data    = ''
 
-    return total_data
+        # Loop until we have all of the data
+        for chunk in response.iter_content(chunk_size = 2**10):
+            total_data    += chunk
+            if total_size:
+                w.status('%s / %s' % (size(total_data), size(total_size)))
+            else:
+                w.status('%s' % size(total_data))
+
+        # Save to the target file if provided
+        if save:
+            if not isinstance(save, (str, unicode)):
+                save = os.path.basename(url)
+                save = save or NamedTemporaryFile(dir='.', delete=False).name
+            with file(save,'wb+') as f:
+                f.write(total_data)
+                w.success('Saved %r (%s)' % (f.name, size(total_data)))
+        else:
+            w.success('%s' % size(total_data))
+
+        return total_data
 

setup.py

@@ -40,7 +40,7 @@ setup(
     url                  = 'https://github.com/Gallopsled/pwntools/',
     download_url         = "https://github.com/Gallopsled/pwntools/tarball/%s" % version,
     install_requires     = ['paramiko','argparse', 'mako', 'pyelftools',
-                            'capstone', 'ropgadget', 'pyserial'],
+                            'capstone', 'ropgadget', 'pyserial', 'requests'],
     license              = "Mostly MIT, some GPL/BSD, see LICENSE-pwntools.txt",
     classifiers          = [
         'Topic :: Security',