draft fix for issue toolswatch#53 (toolswatch#53)

config/constants.py

@@ -10,6 +10,7 @@
 current_dir = os.path.abspath(os.path.dirname(__file__))
 root_dir = os.path.normpath(os.path.join(current_dir, ".."))
 export_dir = os.path.normpath(os.path.join(root_dir, "export"))
+config_dir = os.path.abspath(os.path.expanduser(os.path.join("~", ".vfeed")))
 
 # vFeed Database information
 title = "vFeed - The Correlated Vulnerability and Threat Database"
@@ -23,7 +24,10 @@
 url = "http://www.toolswatch.org/vfeed/"
 db = "vfeed.db"
 db_compressed = "vfeed.db.tgz"
+db_local = os.path.join(config_dir, db)
+db_compressed_local = os.path.join(config_dir, db_compressed)
 update_status = "update"
+update_status_local = os.path.join(config_dir, update_status)
 
 # Third party URLs
 cve_url = "http://cve.mitre.org/cgi-bin/cvename.cgi?name="

lib/common/database.py

@@ -6,7 +6,7 @@
 
 import sys
 import sqlite3
-from config.constants import db
+from config.constants import db_local as db
 from lib.common.utils import check_env
 
 

lib/core/methods/json_dump.py

@@ -6,7 +6,7 @@
 import os
 import json
 import inspect
-from config.constants import title, author, build, repository, twitter, db
+from config.constants import title, author, build, repository, twitter, db_local
 from lib.common.database import Database
 from lib.common.utils import check_env, move_export
 from lib.core.methods import *
@@ -15,7 +15,7 @@
 class ExportJson(object):
     def __init__(self, cve):
         self.cve = cve.upper()
-        self.db = db
+        self.db = db_local
         check_env(self.db)
         (self.cur, self.query) = Database(self.cve).db_init()
         self.data = Database(self.cve, self.cur, self.query).check_cve()

lib/core/methods/patches.py

@@ -13,7 +13,7 @@
 class CvePatches(object):
     def __init__(self, cve):
         self.cve = cve.upper()
-        self.db = db
+        self.db = db_local
         check_env(self.db)
         (self.cur, self.query) = Database(self.cve).db_init()
         self.data = Database(self.cve, self.cur, self.query).check_cve()

lib/core/search.py

@@ -6,7 +6,7 @@
 import json
 import sys
 import re
-from config.constants import db
+from config.constants import db_local as db
 from lib.core.methods import CveExploit
 from lib.common.database import Database
 

lib/core/update.py

@@ -6,19 +6,25 @@
 import sys
 import urllib2
 import tarfile
-from config.constants import db, db_compressed, url, url_test, update_status
+from config.constants import db, db_compressed, url, url_test, update_status, \
+    config_dir, db_local, db_compressed_local, update_status_local
 from lib.common.utils import checksum
 
 
 class Update(object):
     def __init__(self):
         self.db = db
         self.db_compressed = db_compressed
+        self.db_local = db_local
+        self.db_compressed_local = db_compressed_local
+        self.config_dir = config_dir
         self.url_test = url_test
         self.db_url = url
         self.db_update = update_status
+        self.db_update_local = update_status_local
         self.db_download = self.db_url + self.db_compressed
         self.db_status = self.db_url + self.db_update
+        self.db_status_local = os.path.join(config_dir, update_status)
         self.remote_db = self.db_url + self.db_compressed
 
     def update(self):
@@ -30,29 +36,32 @@ def update(self):
         print "[+] Checking connectivity to", self.db_url
         try:
             if urllib2.urlopen(self.url_test):
-                if not os.path.isfile(self.db):
+                if not os.path.isdir(self.config_dir):
+                    os.makedirs(self.config_dir)
+                if not os.path.isfile(self.db_local):
                     print "[+] New install. Downloading the Correlated Vulnerability Database."
-                    self.download(self.remote_db)
+                    self.download(self.remote_db, self.db_compressed_local)
                     print '\n[+] Installing %s ...' % self.db_compressed
                     self.uncompress()
                     self.clean()
                     sys.exit(1)
-                if os.path.isfile(self.db):
+                if os.path.isfile(self.db_local):
                     print "[+] Checking for the latest vFeed Vulnerability Database"
                     self.check_status()
         except urllib2.URLError as e:
             print "[!] Connection error: ", e.reason
             sys.exit()
 
-    def download(self, url):
+    def download(self, url, dest=None):
         """
         This function was found in internet. So thanks to its author wherever he is.
         Just improve it a little by adding the percentage display
         :param url:
         :return:
         """
 
-        self.filename = url.split('/')[-1]
+        self.filename = dest or url.split('/')[-1]
+        self.local = os.path.basename(dest)
         self.u = urllib2.urlopen(url)
         self.f = open(self.filename, 'wb')
         self.meta = self.u.info()
@@ -70,7 +79,7 @@ def download(self, url):
             self.status = r"%10d [%3.0f %%]" % (self.filesize_dl, self.filesize_dl * 100. / self.filesize)
             self.status += chr(8) * (len(self.status) + 1)
             sys.stdout.write("\r[+] Receiving %d out of %s Bytes of %s (%3.0f %%)" % (
-                self.filesize_dl, self.filesize, self.filename, self.filesize_dl * 100. / self.filesize))
+                self.filesize_dl, self.filesize, self.local, self.filesize_dl * 100. / self.filesize))
             sys.stdout.flush()
         self.f.close()
 
@@ -80,29 +89,29 @@ def uncompress(self):
         :return:
         """
 
-        if not os.path.isfile(self.db_compressed):
-            print '[error] ' + self.db_compressed + ' not found'
+        if not os.path.isfile(self.db_compressed_local):
+            print '[error] ' + self.db_compressed_local + ' not found'
             sys.exit()
         try:
-            self.tar = tarfile.open(self.db_compressed, 'r:gz')
-            self.tar.extractall('.')
+            self.tar = tarfile.open(self.db_compressed_local, 'r:gz')
+            self.tar.extractall(self.config_dir)
         except Exception, e:
             print '[error] Database not extracted ', e
 
     def check_status(self):
         """ Check the remote update status and
         update the existing vfeed database if needed
         """
-        self.download(self.db_status)
-        self.hashLocal = checksum(self.db)
-        with open(self.db_update, 'r') as f:
+        self.download(self.db_status, self.db_status_local)
+        self.hashLocal = checksum(self.db_local)
+        with open(self.db_status_local, 'r') as f:
             self.output = f.read()
             self.hashRemote = self.output.split(',')[1]
 
         if self.hashRemote != self.hashLocal:
             print '\n[+] Downloading the recent vFeed Vulnerability Database update'
-            self.download(self.remote_db)
-            print '\n[+] Decompressing %s ' % self.db_compressed
+            self.download(self.remote_db, self.db_compressed_local)
+            print '\n[+] Decompressing %s ' % self.db_compressed_local
             self.uncompress()
 
         if self.hashRemote == self.hashLocal:
@@ -115,9 +124,9 @@ def clean(self):
         """
         print '[+] Cleaning compressed database and update file'
         try:
-            if os.path.isfile(self.db_compressed):
-                os.remove(self.db_compressed)
-            if os.path.isfile(self.db_update):
-                os.remove(self.db_update)
+            if os.path.isfile(self.db_compressed_local):
+                os.remove(self.db_compressed_local)
+            if os.path.isfile(self.db_update_local):
+                os.remove(self.db_update_local)
         except Exception, e:
             print '[!] Already cleaned', e