From 3b9ee1e1e2c19d0a7ee4981b9a06458c676a37d2 Mon Sep 17 00:00:00 2001
From: JayPiKay <jpk@goatpr0n.de>
Date: Sat, 27 Oct 2018 13:15:01 +0200
Subject: [PATCH] Refactoring

---
 requirements.txt |  6 ++++
 sbeam.py         | 84 ++++++++++++++++++++++++++----------------------
 2 files changed, 52 insertions(+), 38 deletions(-)
 create mode 100644 requirements.txt

diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..a71761f
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,6 @@
+certifi==2018.10.15
+chardet==3.0.4
+idna==2.7
+progress==1.4
+requests==2.20.0
+urllib3==1.24
diff --git a/sbeam.py b/sbeam.py
index 9b8c2e0..07e815d 100644
--- a/sbeam.py
+++ b/sbeam.py
@@ -12,48 +12,56 @@
 
 import os
 import requests
+from progress.spinner import Spinner
+from progress.bar import FillingSquaresBar as Bar
 
 
-def get_superlist():
-    with requests.get('http://192.168.178.50:8080/superlist',
-                          headers={'User-Agent': 'sbeam'}) as response:
-        assert response.status_code == 200
-        return response.text.split('\n')
+class SbeamReceiver(object):
+
+    def __init__(self, host):
+        self.host = host
+        self._headers = {'User-Agent': 'sbeam'}
+        self.url = f'http://{host}:8080'
+        self.superlist()
+
+    def superlist(self):
+        url = f'{self.url}/superlist'
+        with requests.get(f'{self.url}/superlist',
+                          headers=self._headers) as response:
+            assert response.status_code == 200
+            self.filelist = response.text.split('\n')
+            try:
+                self._total_bytes = 0
+                for item in self.filelist:
+                    self._total_bytes += int(item.split('||')[1])
+            finally:
+                return len(self.filelist) - 1
+
+    def save_stream(self, filename, chunk_size=16392):
+        with requests.get(f'{self.url}/getstream', stream=True,
+                          headers=self._headers) as response:
+            assert response.status_code == 200
+            with open(filename, 'wb') as handle:
+                for chunk in Bar('Downloading',
+                                 max=self._total_bytes/chunk_size) \
+                        .iter(response.iter_content(chunk_size=chunk_size)):
+                    handle.write(chunk)
 
 
-def save_stream(filename):
-    with requests.get('http://192.168.178.50:8080/getstream', stream=True,
-                          headers={'User-Agent': 'sbeam'}) as response:
-        assert response.status_code == 200
-        with open(filename, 'wb') as handle:
-            # length = response.headers.get('Content-length')
-            # print(length)
-            for chunk in response.iter_content(chunk_size=8196):
-                handle.write(chunk)
-
-
-def split_stream(data, directory, superlist):
-    with open(data, 'rb') as blob:
-        for item in superlist:
-            if not item:
-                break
-            filename, filesize, filepath, timestamp = item.split('||')
-            fullpath = os.path.join(directory, filepath[1:])
-            os.makedirs(fullpath, exist_ok=True)
-            print(f'Saving {filename} ({filesize} bytes) to disk...')
-            with open(os.path.join(fullpath, filename), 'wb') as handle:
-                handle.write(blob.read(int(filesize)))
+    def split_stream(self, data, directory):
+        with open(data, 'rb') as blob:
+            for item in self.filelist:
+                if not item:
+                    break
+                filename, filesize, filepath, timestamp = item.split('||')
+                fullpath = os.path.join(directory, filepath[1:])
+                os.makedirs(fullpath, exist_ok=True)
+                print(f'Unpacking {filename} ({filesize} bytes)...')
+                with open(os.path.join(fullpath, filename), 'wb') as handle:
+                    handle.write(blob.read(int(filesize)))
 
 
 if __name__ == '__main__':
-    to_download = 0
-    superlist = get_superlist()
-    try:
-        for item in superlist:
-            to_download += int(item.split('||')[1])
-    except IndexError:
-        print(f'Download stream contains {len(superlist) - 1} items.')
-    finally:
-        print(f'Target download size is {to_download} bytes.')
-        save_stream('data.gz')
-        split_stream('data.gz', 'received', superlist)
+    sbeam_rec = SbeamReceiver('192.168.178.50')
+    sbeam_rec.save_stream('data.gz')
+    sbeam_rec.split_stream('data.gz', 'received')