From ff29fe37e481f24326d6c154b752b041f9c3c99b Mon Sep 17 00:00:00 2001
From: Mark <lo127001@gmail.com>
Date: Wed, 20 Jun 2018 10:40:06 -0400
Subject: [PATCH] fixed update_diskcache and use pprint when printing the
 entire whois record

---
 README.md           |  1 +
 domain_stats.py     |  4 +++-
 update_diskcache.py | 10 +++++-----
 3 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index 0cae900..63b2e13 100644
--- a/README.md
+++ b/README.md
@@ -145,6 +145,7 @@ student@SEC573:~$ python domain_stats.py --preload 100 -a ~/Downloads/top-1m.csv
 ---
 
 CHANGE LOG:  Version 1.0 -> 1.1
+
 *Update code to work in either Python2 or Python3
 *Changed Default Content type of server response to TEXT.
 *Added -d option which will load top 1000 domains from disk when -a top-1m.csv and --preload are not used.
diff --git a/domain_stats.py b/domain_stats.py
index c3e7c39..b056bcd 100644
--- a/domain_stats.py
+++ b/domain_stats.py
@@ -25,6 +25,7 @@
 import os
 import datetime
 import pickle
+from pprint import pformat
 
 try:
     import whois
@@ -119,7 +120,8 @@ def do_GET(self):
                 finally:
                     self.server.cache_lock.release()
             if not fields:
-                self.wfile.write(str(domain_info).encode("latin-1"))
+                dinfo = pformat(domain_info)
+                self.wfile.write(dinfo.encode("latin-1"))
             else:
                 if self.server.args.verbose: self.server.safe_print("processing fields %s" % (fields))
                 if domain_info.get('status','') == "NOT FOUND":
diff --git a/update_diskcache.py b/update_diskcache.py
index 8f3aa38..01af4bb 100644
--- a/update_diskcache.py
+++ b/update_diskcache.py
@@ -46,17 +46,16 @@ def preload_domains(domain_list, delay=0.1):
 parser.add_argument('-f','--file',help='Name of the file to write.',default="domain_cache.dst")
 parser.add_argument('-a','--append',action="store_true", required=False,help='Append to existing file instead of overwriting.')
 args = parser.parse_args()
-print args
 
 cache = {}
 try:
-    alexa = open("top-1m.csv").readlines()[900:args.count]
+    alexa = open("top-1m.csv").readlines()[:args.count]
 except Exception as e:
     raise(Exception("Cant find your alexa top-1m.csv file. {0}".format(str(e))))
 
 if args.append:
     try:
-        fh = open("domain_cache.dst") 
+        fh = open("domain_cache.dst","rb") 
         cache = pickle.load(fh)
         fh.close()
     except Exception as e:
@@ -65,8 +64,9 @@ def preload_domains(domain_list, delay=0.1):
 preload_domains(alexa)
 
 try:
-    fh = open(args.file,"w")
-    pickle.dump(cache, fh, protocol=2)
+    fh = open(args.file,"wb")
+    fh.write(pickle.dumps(cache,protocol=2))
+    fh.flush()
     fh.close()
 except Exception as e:
     raise(Exception("Unable to create your disk cache file. {0}".format(str(e))))