Merge pull request #65 from P-T-I/CveXplore-63

Add schema version into CveXplore

CveXplore/.schema_version

@@ -0,0 +1,4 @@
+{
+    "version": "1.3",
+    "rebuild_needed": true
+}

CveXplore/cli_cmds/db_cmds/commands.py

@@ -1,5 +1,8 @@
 import click
 
+from CveXplore.cli_cmds.cli_utils.utils import printer
+from CveXplore.database.maintenance.Config import Configuration
+
 
 @click.group(
     "database", invoke_without_command=True, help="Database update / populate commands"
@@ -19,3 +22,22 @@ def update_cmd(ctx):
 @click.pass_context
 def initialize_cmd(ctx):
     ctx.obj["data_source"].database.initialize()
+
+
+@db_cmd.group("sources", invoke_without_command=True, help="Database source management")
+@click.option("--pretty", is_flag=True, help="Pretty print the output")
+@click.pass_context
+def sources_cmd(ctx, pretty):
+
+    config = Configuration()
+
+    if ctx.invoked_subcommand is None:
+        printer(
+            input_data=config.SOURCES,
+            pretty=pretty,
+        )
+    else:
+        printer(
+            input_data=config.SOURCES,
+            pretty=pretty,
+        )

CveXplore/database/maintenance/Config.py

@@ -3,6 +3,7 @@
 =============
 """
 import datetime
+import json
 import os
 import re
 
@@ -16,16 +17,16 @@ class Configuration(object):
 
     CVE_START_YEAR = os.getenv("CVE_START_YEAR", 2002)
 
-    SOURCES = os.getenv(
-        "SOURCES",
-        {
+    if os.getenv("SOURCES") is not None:
+        SOURCES = json.loads(os.getenv("SOURCES"))
+    else:
+        SOURCES = {
             "cve": "https://nvd.nist.gov/feeds/json/cve/1.1/",
             "cpe": "https://nvd.nist.gov/feeds/json/cpematch/1.0/nvdcpematch-1.0.json.zip",
             "cwe": "https://cwe.mitre.org/data/xml/cwec_v4.4.xml.zip",
-            "capec": "https://capec.mitre.org/data/xml/capec_v3.4.xml",
+            "capec": "https://capec.mitre.org/data/xml/capec_v3.5.xml",
             "via4": "https://www.cve-search.org/feeds/via4.json",
-        },
-    )
+        }
 
     HTTP_PROXY = os.getenv("HTTP_PROXY", "")
 

CveXplore/database/maintenance/DatabaseSchemaChecker.py

@@ -0,0 +1,53 @@
+import json
+import logging
+import os
+
+from CveXplore.database.connection.mongo_db import MongoDBConnection
+
+runPath = os.path.dirname(os.path.realpath(__file__))
+
+
+class SchemaChecker(object):
+    def __init__(self):
+        with open(os.path.join(runPath, "../../.schema_version")) as f:
+            self.schema_version = json.loads(f.read())
+
+        database = MongoDBConnection(**json.loads(os.getenv("MONGODB_CON_DETAILS")))
+
+        self.dbh = database._dbclient["schema"]
+
+        self.logger = logging.getLogger("SchemaChecker")
+
+    def create_indexes(self):
+        # hack for db_updater.py to put this class in the posts variable and run the update method
+        self.logger.info("Updating schema version")
+        self.update()
+        self.logger.info("Update schema version done!")
+
+    def update(self):
+        try:
+
+            current_record = list(self.dbh.find({}))
+
+            if len(current_record) != 0:
+                current_record[0]["version"] = self.schema_version["version"]
+                current_record[0]["rebuild_needed"] = self.schema_version[
+                    "rebuild_needed"
+                ]
+
+                self.dbh.update_one(
+                    {"_id": current_record[0]["_id"]}, {"$set": current_record[0]}
+                )
+            else:
+                current_record = {
+                    "version": self.schema_version["version"],
+                    "rebuild_needed": self.schema_version["rebuild_needed"],
+                }
+
+                self.dbh.insert_one(current_record)
+        except AttributeError:
+            current_record = {
+                "version": self.schema_version["version"],
+                "rebuild_needed": self.schema_version["rebuild_needed"],
+            }
+            self.dbh.insert_one(current_record)

CveXplore/database/maintenance/main_updater.py

@@ -2,6 +2,7 @@
 Main Updater
 ============
 """
+from CveXplore.database.maintenance.DatabaseSchemaChecker import SchemaChecker
 from CveXplore.database.maintenance.Sources_process import (
     CPEDownloads,
     CVEDownloads,
@@ -35,7 +36,10 @@ def __init__(self, datasource):
             {"name": "via4", "updater": VIADownloads},
         ]
 
-        self.posts = [{"name": "ensureindex", "updater": DatabaseIndexer}]
+        self.posts = [
+            {"name": "ensureindex", "updater": DatabaseIndexer},
+            {"name": "schema", "updater": SchemaChecker},
+        ]
 
     def update(self):
         """

setup.py

@@ -36,7 +36,7 @@
     description="Package for interacting with cve-search",
     long_description=README,
     long_description_content_type="text/x-rst",
-    package_data={"CveXplore": ["LICENSE", "VERSION", ".cvexplore-complete.bash"]},
+    package_data={"CveXplore": ["LICENSE", "VERSION", ".cvexplore-complete.bash", ".schema_version"]},
     entry_points='''
             [console_scripts]
             cvexplore=CveXplore.cli:main