From 92ffd1b7dadb33410569788c85b55c1e5f903862 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rapha=C3=ABl=20Vinot?= <raphael@vinot.info>
Date: Tue, 6 Feb 2018 01:23:00 +0100
Subject: [PATCH] new: Allow to download the bodies

---
 README.md                    |  1 +
 lookyloo/__init__.py         | 21 +++++++++++++++++++--
 lookyloo/static/tree.js      | 33 ++++++++++++++++++++++++++-------
 lookyloo/templates/main.html |  1 +
 4 files changed, 47 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index f31810e..420ac01 100644
--- a/README.md
+++ b/README.md
@@ -49,6 +49,7 @@ sudo docker run -p 8050:8050 -p 5023:5023 scrapinghub/splash --disable-ui --disa
 pip install -r requirements.txt
 pip install -e .
 wget https://d3js.org/d3.v4.min.js -O lookyloo/static/d3.v4.min.js
+wget https://cdn.rawgit.com/eligrey/FileSaver.js/5733e40e5af936eb3f48554cf6a8a7075d71d18a/FileSaver.js -O lookyloo/static/FileSaver.js
 ```
 # Run the app locally
 
diff --git a/lookyloo/__init__.py b/lookyloo/__init__.py
index 99057ad..f446a4c 100644
--- a/lookyloo/__init__.py
+++ b/lookyloo/__init__.py
@@ -6,7 +6,7 @@ import json
 from har2tree import CrawledTree
 from scrapysplashwrapper import crawl
 
-from flask import Flask, render_template, request, session
+from flask import Flask, render_template, request, session, send_file
 from flask_bootstrap import Bootstrap
 
 from glob import glob
@@ -15,6 +15,10 @@ from datetime import datetime
 
 import pickle
 import tempfile
+import pathlib
+
+from zipfile import ZipFile, ZIP_DEFLATED
+from io import BytesIO
 
 app = Flask(__name__)
 
@@ -30,6 +34,8 @@ app.debug = True
 HAR_DIR = 'scraped'
 SPLASH = 'http://127.0.0.1:8050'
 
+pathlib.Path(HAR_DIR).mkdir(parents=True, exist_ok=True)
+
 
 @app.before_request
 def session_management():
@@ -103,7 +109,18 @@ def urlnode_details(node_uuid):
     with open(session["tree"], 'rb') as f:
         ct = pickle.load(f)
     urlnode = ct.root_hartree.get_url_node_by_uuid(node_uuid)
-    return urlnode.to_json()
+
+    to_return = BytesIO()
+    if hasattr(urlnode, 'body'):
+        with ZipFile(to_return, 'a', ZIP_DEFLATED, False) as zfile:
+            zfile.writestr(urlnode.filename, urlnode.body.getvalue())
+        to_return.seek(0)
+    # return send_file(urlnode.body, mimetype='application/zip',
+    #                 as_attachment=True, attachment_filename='file.zip')
+    with open('foo.bin', 'wb') as f:
+        f.write(to_return.getvalue())
+    return send_file(to_return, mimetype='application/zip',
+                     as_attachment=True, attachment_filename='file.zip')
 
 
 @app.route('/tree/<int:tree_id>', methods=['GET'])
diff --git a/lookyloo/static/tree.js b/lookyloo/static/tree.js
index 913d18b..314b93d 100644
--- a/lookyloo/static/tree.js
+++ b/lookyloo/static/tree.js
@@ -68,7 +68,7 @@ function collapse(d) {
     d._children.forEach(collapse)
     d.children = null
   }
-}
+};
 
 // Gets the size of the box and increase it
 function getBB(selection) {
@@ -76,15 +76,30 @@ function getBB(selection) {
         d.data.total_width = d.data.total_width ? d.data.total_width : 0;
         d.data.total_width += this.getBBox().width;
     })
+};
+
+function str2bytes (str) {
+    var bytes = new Uint8Array(str.length);
+    for (var i=0; i<str.length; i++) {
+        bytes[i] = str.charCodeAt(i);
+    }
+    return bytes;
 }
 
 function urlnode_click(d) {
     var url = "url/" + d.data.uuid;
-    d3.json(url, function(error, u) {
-        if (error) throw error;
-        console.log(u)
-    })
-}
+	var xhr = new XMLHttpRequest();
+	xhr.open('GET', url, true);
+	xhr.responseType = "blob";
+	xhr.withCredentials = true;
+	xhr.onreadystatechange = function (){
+		if (xhr.readyState === 4) {
+			var blob = xhr.response;
+			saveAs(blob, 'file.zip');
+		}
+	};
+	xhr.send();
+};
 
 // What happen when clicking on a domain (load a modal display)
 function hostnode_click(d) {
@@ -208,7 +223,11 @@ function text_entry(parent_svg, relative_x_pos, relative_y_pos, onclick_callback
           .style("opacity", .9)
           .text(function(d) {
               d.data.total_width = 0; // reset total_width
-              return d.data.name;
+              to_display = d.data.name
+              if (d.data.urls_count) {
+                  to_display += ' (' + d.data.urls_count + ')';
+              };
+              return to_display;
           })
           .on('click', onclick_callback);
 
diff --git a/lookyloo/templates/main.html b/lookyloo/templates/main.html
index d8522a1..99be2ff 100644
--- a/lookyloo/templates/main.html
+++ b/lookyloo/templates/main.html
@@ -2,6 +2,7 @@
 
 {% block scripts %}
     {{ super() }}
+    <script src='{{ url_for('static', filename='FileSaver.js') }}'></script>
     <script src='{{ url_for('static', filename='d3.v4.min.js') }}'></script>
 {% endblock %}