From 9a3abe91381d16a9da8d909a504f890be9c9fb76 Mon Sep 17 00:00:00 2001 From: qtfkwk Date: Tue, 16 Feb 2016 15:15:28 -0700 Subject: [PATCH] draft fix for issue #53 (https://github.com/toolswatch/vFeed/issues/53) --- config/constants.py | 4 +++ lib/common/database.py | 2 +- lib/core/methods/json_dump.py | 4 +-- lib/core/methods/patches.py | 2 +- lib/core/search.py | 2 +- lib/core/update.py | 49 +++++++++++++++++++++-------------- 6 files changed, 38 insertions(+), 25 deletions(-) diff --git a/config/constants.py b/config/constants.py index 6ac5d46..cf17c32 100644 --- a/config/constants.py +++ b/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=" diff --git a/lib/common/database.py b/lib/common/database.py index ae022a3..632ca80 100644 --- a/lib/common/database.py +++ b/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 diff --git a/lib/core/methods/json_dump.py b/lib/core/methods/json_dump.py index 2cd161c..bb4b78d 100644 --- a/lib/core/methods/json_dump.py +++ b/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() diff --git a/lib/core/methods/patches.py b/lib/core/methods/patches.py index c4695d3..aa0d86d 100644 --- a/lib/core/methods/patches.py +++ b/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() diff --git a/lib/core/search.py b/lib/core/search.py index a24a88f..9a677de 100644 --- a/lib/core/search.py +++ b/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 diff --git a/lib/core/update.py b/lib/core/update.py index 933a7bf..a27736f 100644 --- a/lib/core/update.py +++ b/lib/core/update.py @@ -6,7 +6,8 @@ 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 @@ -14,11 +15,16 @@ 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,21 +36,23 @@ 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 @@ -52,7 +60,8 @@ def download(self, 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,12 +89,12 @@ 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 @@ -93,16 +102,16 @@ 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