diff --git a/README.md b/README.md index 640d86a..dc42b18 100644 --- a/README.md +++ b/README.md @@ -136,5 +136,5 @@ cpe (vendor:product) per version to give a probability of the CPE appearance. Software is open source and released under a 2-Clause BSD License -Copyright (C) 2021 Alexandre Dulaunoy -Copyright (C) 2021 Esa Jokinen +Copyright (C) 2021 Alexandre Dulaunoy +Copyright (C) 2021 Esa Jokinen diff --git a/bin/import.py b/bin/import.py index 398dd8d..87fb808 100644 --- a/bin/import.py +++ b/bin/import.py @@ -9,6 +9,7 @@ import gzip import shutil import xml.sax import redis +import time # Configuration cpe_path = '../data/official-cpe-dictionary_v2.3.xml' @@ -25,6 +26,7 @@ class CPEHandler( xml.sax.ContentHandler ): self.refs = [] self.itemcount = 0 self.wordcount = 0 + self.start_time = time.time() def startElement(self, tag, attributes): self.CurrentData = tag @@ -38,7 +40,7 @@ class CPEHandler( xml.sax.ContentHandler ): def characters(self, data): if self.title_seen: self.title = self.title + data - + def endElement(self, tag): if tag == 'title': self.record['title'] = self.title @@ -58,7 +60,8 @@ class CPEHandler( xml.sax.ContentHandler ): self.record = {} self.itemcount += 1 if self.itemcount % 5000 == 0: - print ("... {} items processed ({} words)".format(str(self.itemcount), str(self.wordcount))) + time_elapsed = round( time.time() - self.start_time ) + print (f"... {self.itemcount} items processed ({self.wordcount} words) in {time_elapsed} seconds") def CPEExtractor( cpe=None ): @@ -70,9 +73,9 @@ def CPEExtractor( cpe=None ): record['product'] = cpefield[4] cpeline = "" for cpeentry in cpefield[:5]: - cpeline = "{}:{}".format(cpeline, cpeentry) - record['cpeline'] = cpeline[1:] - return record + cpeline = f"{cpeline}:{cpeentry}" + record['cpeline'] = cpeline[1:] + return record def canonize( value=None ): value = value.lower() @@ -82,9 +85,9 @@ def canonize( value=None ): def insert( word=None, cpe=None): if cpe is None or word is None: return False - rdb.sadd('w:{}'.format(word), cpe) - rdb.zadd('s:{}'.format(word), {cpe: 1}, incr=True) - rdb.zadd('rank:cpe', {cpe: 1}, incr=True) + rdb.sadd(f"w:{word}", cpe) + rdb.zadd(f"s:{word}", {cpe: 1}, incr=True) + rdb.zadd("rank:cpe", {cpe: 1}, incr=True) if __name__ == '__main__': @@ -94,33 +97,33 @@ if __name__ == '__main__': args = argparser.parse_args() if args.replace == 0 and rdb.dbsize() > 0: - print("Warning! The Redis database already has " + str(rdb.dbsize()) + " keys.") + print(f"Warning! The Redis database already has {rdb.dbsize()} keys.") print("Use --replace if you want to flush the database and repopulate it.") sys.exit(1) if args.download > 0 or not os.path.isfile(cpe_path): - print("Downloading CPE data from " + cpe_source + " ...") + print(f"Downloading CPE data from {cpe_source} ...") try: - urllib.request.urlretrieve(cpe_source, cpe_path + ".gz") + urllib.request.urlretrieve(cpe_source, f"{cpe_path}.gz") except (urllib.error.HTTPError, urllib.error.URLError, FileNotFoundError, PermissionError) as e: print(e) sys.exit(1) - print("Uncompressing {}.gz ...".format(cpe_path)) + print(f"Uncompressing {cpe_path}.gz ...") try: - with gzip.open(cpe_path + ".gz", 'rb') as cpe_gz: + with gzip.open(f"{cpe_path}.gz", 'rb') as cpe_gz: with open(cpe_path, 'wb') as cpe_xml: shutil.copyfileobj(cpe_gz, cpe_xml) - os.remove(cpe_path + ".gz") + os.remove(f"{cpe_path}.gz") except (FileNotFoundError, PermissionError) as e: print(e) sys.exit(1) elif os.path.isfile(cpe_path): - print("Using existing file {} ...".format(cpe_path)) + print(f"Using existing file {cpe_path} ...") if rdb.dbsize() > 0: - print("Flushing {} keys from the database...".format(str(rdb.dbsize()))) + print(f"Flushing {rdb.dbsize()} keys from the database...") rdb.flushdb() print("Populating the database (please be patient)...") @@ -128,4 +131,4 @@ if __name__ == '__main__': Handler = CPEHandler() parser.setContentHandler( Handler ) parser.parse(cpe_path) - print("Done! {} keys inserted.".format(str(rdb.dbsize()))) + print(f"Done! {rdb.dbsize()} keys inserted.") diff --git a/bin/server.py b/bin/server.py index 62f01cd..d256c6a 100644 --- a/bin/server.py +++ b/bin/server.py @@ -5,17 +5,17 @@ import os import sys import falcon from wsgiref.simple_server import make_server -import requests -from datetime import datetime import json +# Configuration +port = 8000 + runPath = os.path.dirname(os.path.realpath(__file__)) sys.path.append(os.path.join(runPath, "..")) from lib.cpeguesser import CPEGuesser class Search(): def on_post(self, req, resp): - ret = [] data_post = req.bounded_stream.read() js = data_post.decode('utf-8') try: @@ -23,10 +23,10 @@ class Search(): except ValueError: resp.status = falcon.HTTP_400 resp.media = "Missing query array or incorrect JSON format" - return + return if 'query' in q: - pass + pass else: resp.status = falcon.HTTP_400 resp.media = "Missing query array or incorrect JSON format" @@ -39,6 +39,12 @@ if __name__ == '__main__': app = falcon.App() app.add_route('/search', Search()) - with make_server('', 8000, app) as httpd: - print('Serving on port 8000...') - httpd.serve_forever() + try: + with make_server('', port, app) as httpd: + print(f"Serving on port {port}...") + httpd.serve_forever() + except OSError as e: + print (e) + sys.exit(1) + except KeyboardInterrupt: + sys.exit(0) diff --git a/lib/cpeguesser.py b/lib/cpeguesser.py index 03c7783..aa32b0d 100644 --- a/lib/cpeguesser.py +++ b/lib/cpeguesser.py @@ -10,7 +10,7 @@ class CPEGuesser(): def guessCpe(self, words): k=[] for keyword in words: - k.append('w:{}'.format(keyword.lower())) + k.append(f"w:{keyword.lower()}") maxinter = len(k) cpes = []