dtube: cache support

1 files changed, 12 insertions, 3 deletions

diff --git a/dtube/dtube.py b/dtube/dtube.py
index 7f5c821..d1af0ff 100755
--- a/dtube/dtube.py
+++ b/dtube/dtube.py
@@ -4,12 +4,14 @@
 # 2013 Laurent Ghigonis <laurent@p1sec.com>
 #
 # TODO
-# * command line usage
 # * cache
+# * ip dns country
+# * -t
 
 import sys
 import os
 import requests
+import requests_cache # pip install requests_cache
 import pprint
 import time
 import random
@@ -87,7 +89,8 @@ class He:
           raise Exception("Target not found")
         self.html.raise_for_status()
       self.tree = etree.HTML(self.html.content)
-      time.sleep(random.random() * 3) # Forced to be kind
+      if not getattr(self.html, 'from_cache', False):
+        time.sleep(random.random() * 3) # Forced to be kind
 
 class HeAS(He):
   url_prefix = '/AS'
@@ -219,12 +222,18 @@ if __name__ == '__main__':
   #                        help='Query type: range, A, PTR, ip, country')
   #parser.add_option('-f', action="store", dest="filter", default=None, type=str,
   #                        help='Filter regex for company name')
+  parser.add_option('-n', action="store_true", dest="url_nocache", default=False,
+                          help='Do not cache answers')
   parser.add_option('-v', action="store_true", dest="verbose", default=False,
-                          help='verbose')
+                          help='Verbose')
   (options, args) = parser.parse_args()
   if (len(args) != 1):
     print usage()
     sys.exit(1)
   target = args[0]
+
+  if options.url_nocache is False:
+    requests_cache.install_cache('cache_dtube')
+
   he = He.initTarget(target, verbose=options.verbose)
   he.show()