Various fixes and improvements

- Refactoring of code
- Enhance README for mariadb
- Improve protection against command and SQL injection
    - use prepared statements where possible
    - validate database and table names
    - use lists for calling subprocesses
- Do not build test cases for mariadb connector
- Use sqlite as default webserver again
- Remove 'CREATE_SQL_STATEMENTS_FILE' as config attribute
- Add back code that was accidentally deleted in create_db
- Update test cases to match with changes in cpe_search
- Do not remove MariaDB backup file if MariaDB is down

README.md

@@ -89,15 +89,6 @@ It is also possible to run a web server that provides this tool's functionality
 ```bash
 ./web_server.py
 ```
-The web server uses mariadb as database. To change it to sqlite, simply change the CONFIG_FILE variable at top of the web_server.py script to `config.json`. 
-It is recommend to change the following values in `/etc/my.cnf` to make some useful changes:
-```
-query_cache_type = 1
-query_cache_size = 192M
-innodb_buffer_pool_size = 8G
-thread_handling=pool-of-threads
-```
-innodb_buffer_pool_size should be set to approximately 80% of available memory (https://mariadb.com/kb/en/innodb-system-variables/#innodb_buffer_pool_size)
 ```bash
 gunicorn --worker-class=gevent --worker-connections=50 --workers=3 --bind '0.0.0.0:8000' wsgi:app
 ```
@@ -108,5 +99,22 @@ Finally, you can also use Nginx as a reverse proxy. A sample configuration file
 gunicorn --worker-class=gevent --worker-connections=50 --workers=3 --bind 'unix:/tmp/gunicorn.sock' wsgi:app
 ```
 
+## MariaDB as Second Database Option
+As alternative to the preconfigured SQLite, you can use *MariaDB* as database. A sample configuration file for MariaDB is provided in [``config_mariadb.json``](https://github.com/ra1nb0rn/search_vulns/blob/master/config_mariadb.json).
+
+Make sure that you adjust the values for MariaDB in the configuration file to your MariaDB configuration (*user*, *password*, *host*, *port*).
+
+To use MariaDB instead of *SQLite* for the webserver, simply change the CONFIG_FILE variable in ``web_server.py`` to your config file (e.g. ``config_mariadb.json``).
+It is recommend to change the following values in ``/etc/my.cnf`` to improve the performance of MariaDB:
+```
+[mariadb]
+
+query_cache_type = 1
+query_cache_size = 192M
+innodb_buffer_pool_size = 8G
+thread_handling = pool-of-threads
+```
+`innodb_buffer_pool_size` should be set to approximately 80% of available memory (see [the official documentation](https://mariadb.com/kb/en/innodb-system-variables/#innodb_buffer_pool_size)).
+
 ## License
 *search_vulns* is licensed under the MIT license, see [here](https://github.com/ra1nb0rn/search_vulns/blob/master/LICENSE).

config.json

@@ -2,12 +2,10 @@
     "DATABASE_NAME": "vulndb.db3",
     "MAN_EQUIVALENT_CPES_FILE": "man_equiv_cpes.json",
     "CVE_EDB_MAP_FILE": "cveid_to_edbid.json",
-    "CREATE_SQL_STATEMENTS_FILE": "create_sql_statements.json",
     "NVD_API_KEY": "",
     "cpe_search": {
         "DATABASE_NAME": "cpe_search/cpe-search-dictionary.db3",
         "DEPRECATED_CPES_FILE": "cpe_search/deprecated-cpes.json",
-        "CREATE_SQL_STATEMENTS_FILE": "cpe_search/create_sql_statements.json",
         "NVD_API_KEY": ""
     },
     "DATABASE": {

config_mariadb.json

@@ -2,12 +2,10 @@
     "DATABASE_NAME": "vulndb",
     "MAN_EQUIVALENT_CPES_FILE": "man_equiv_cpes.json",
     "CVE_EDB_MAP_FILE": "cveid_to_edbid.json",
-    "CREATE_SQL_STATEMENTS_FILE": "create_sql_statements.json",
     "NVD_API_KEY": "",
     "cpe_search": {
         "DATABASE_NAME": "cpe_search_dictionary",
         "DEPRECATED_CPES_FILE": "cpe_search/deprecated-cpes.json",
-        "CREATE_SQL_STATEMENTS_FILE": "cpe_search/create_sql_statements.json",
         "NVD_API_KEY": ""
     },
     "DATABASE": {

db_creation_src/CMakeLists.txt

@@ -9,9 +9,9 @@ include_directories(SQLiteCpp/include)
 add_subdirectory(SQLiteCpp)
 
 # Include mariadb-connector-cpp library
-set(CONC_WITH_UNIT_TESTS "Off")
+set(CONC_WITH_UNIT_TESTS OFF)
 set(CMAKE_BUILD_TYPE "RelWithDebInfo")
-set(WITH_UNIT_TESTS "Off")
+set(WITH_UNIT_TESTS OFF CACHE INTERNAL "")
 include_directories(mariadb-connector-cpp/include)
 # workaround until mariadb fix issue in test/CMakeLists.txt
 include_directories("${CMAKE_BINARY_DIR}/mariadb-connector-cpp/test")

db_creation_src/create_db.cpp

@@ -11,6 +11,7 @@
 #include <climits>
 #include <unordered_set>
 #include <unordered_map>
+#include <regex>
 #include "database_wrapper.h"
 #include "prepared_statement.h"
 
@@ -61,6 +62,18 @@ void handle_exception(T &e) {
     }
 }
 
+bool is_safe_database_name(std::string dbName) {
+    // Check if database name contains any special characters or keywords
+    std::regex pattern("[^a-zA-Z0-9_-]");
+
+    if (std::regex_search(dbName, pattern)) {
+        return false; // Database name is malicious
+    }
+
+    return true; // Database name is safe
+}
+
+
 int add_to_db(DatabaseWrapper *db, const std::string &filepath) {
     // Begin transaction
     db->start_transaction();
@@ -253,7 +266,7 @@ int add_to_db(DatabaseWrapper *db, const std::string &filepath) {
                     else
                         cve_cpe_query->bind(6, false);
 
-                    try{
+                    try {
                         cve_cpe_query->execute();
                     }
                     catch (SQLite::Exception& e) {
@@ -314,9 +327,15 @@ int main(int argc, char *argv[]) {
 
     std::unique_ptr<DatabaseWrapper> db;
 
+    // validate given database name
+    if (database_type != "sqlite" && !is_safe_database_name(config["DATABASE_NAME"])) {
+        std::cout << "Potentially malicious database name detected. Abort creation of database" << std::endl;
+        return EXIT_FAILURE;
+    }
+
     try{
         // create database connection
-        if ( database_type == "sqlite")
+        if (database_type == "sqlite")
             db = std::make_unique<SQLiteDB>(outfile);
         else{
             db = std::make_unique<MariaDB>(config);
@@ -360,4 +379,15 @@ int main(int argc, char *argv[]) {
     }
     // close database connection
     db->close_connection();
+
+    // print duration of building process
+    auto time = std::chrono::high_resolution_clock::now() - start_time;
+
+    char *db_abs_path = realpath(outfile.c_str(), NULL);
+    std::cout << "Database creation took " <<
+    (float) (std::chrono::duration_cast<std::chrono::microseconds>(time).count()) / (1e6) << "s .\n";
+    std::cout << "Local copy of NVD created as " << db_abs_path << " ." << std::endl;
+    free(db_abs_path);
+    return EXIT_SUCCESS;
+
 }

db_creation_src/database_wrapper.cpp

@@ -31,10 +31,10 @@ void SQLiteDB::execute_query(std::string query) {
 
 void SQLiteDB::create_prepared_statements() {
     // init prepared statements
-    cve_query =  std::make_unique<SQLitePreparedStatement>(*db, CVE_QUERY_FRAGMENT);
-    cve_cpe_query =  std::make_unique<SQLitePreparedStatement>(*db, CVE_CPE_QUERY_FRAGMENT);
-    add_exploit_ref_query =  std::make_unique<SQLitePreparedStatement>(*db, NVD_EXPLOIT_REFS_FRAGMENT);
-    add_cveid_exploit_ref_query =  std::make_unique<SQLitePreparedStatement>(*db, CVE_NVD_EXPLOITS_REFS_FRAGMENT);
+    cve_query = std::make_unique<SQLitePreparedStatement>(*db, CVE_QUERY_FRAGMENT);
+    cve_cpe_query = std::make_unique<SQLitePreparedStatement>(*db, CVE_CPE_QUERY_FRAGMENT);
+    add_exploit_ref_query = std::make_unique<SQLitePreparedStatement>(*db, NVD_EXPLOIT_REFS_FRAGMENT);
+    add_cveid_exploit_ref_query = std::make_unique<SQLitePreparedStatement>(*db, CVE_NVD_EXPLOITS_REFS_FRAGMENT);
 }
 
 void SQLiteDB::commit() {
@@ -63,7 +63,8 @@ MariaDB::MariaDB(nlohmann::json config) {
     sql::Properties properties({{"user", user}, {"password", password}});
 
     conn = std::unique_ptr<sql::Connection>(driver->connect(url, properties));
-
+
+    // set up database
     std::string create_db_query = "CREATE OR REPLACE DATABASE "+database+";";
     std::unique_ptr<sql::Statement> stmnt(conn->createStatement());
     stmnt->executeQuery(create_db_query);
@@ -77,10 +78,10 @@ void MariaDB::execute_query(std::string query) {
 
 void MariaDB::create_prepared_statements() {
     // init prepared statements
-    cve_query =  std::make_unique<MariaDBPreparedStatement>(conn, CVE_QUERY_FRAGMENT);
-    cve_cpe_query =  std::make_unique<MariaDBPreparedStatement>(conn, CVE_CPE_QUERY_FRAGMENT);
-    add_exploit_ref_query =  std::make_unique<MariaDBPreparedStatement>(conn, NVD_EXPLOIT_REFS_FRAGMENT);
-    add_cveid_exploit_ref_query =  std::make_unique<MariaDBPreparedStatement>(conn, CVE_NVD_EXPLOITS_REFS_FRAGMENT);
+    cve_query = std::make_unique<MariaDBPreparedStatement>(conn, CVE_QUERY_FRAGMENT);
+    cve_cpe_query = std::make_unique<MariaDBPreparedStatement>(conn, CVE_CPE_QUERY_FRAGMENT);
+    add_exploit_ref_query = std::make_unique<MariaDBPreparedStatement>(conn, NVD_EXPLOIT_REFS_FRAGMENT);
+    add_cveid_exploit_ref_query = std::make_unique<MariaDBPreparedStatement>(conn, CVE_NVD_EXPLOITS_REFS_FRAGMENT);
 }
 
 void MariaDB::commit() {

export_database_as_csv.py

@@ -1,12 +1,17 @@
 import sqlite3
 try: # only use mariadb module if installed
     import mariadb
-except:
+except ImportError:
     pass
 import csv
 import sys
 import os
 
+
+def is_safe_table_name(table_name):
+    return all([c.isalnum() or c in ('-', '_') for c in table_name])
+
+
 def export_tables_to_csv(database_file):
     # Connect to the SQLite database
     conn = sqlite3.connect(database_file)
@@ -19,8 +24,12 @@ def export_tables_to_csv(database_file):
     # Export each table to a separate CSV file
     for table in tables:
         table_name = table[0]
-        csv_file = os.path.join(os.path.dirname(os.path.realpath(__file__)), f"{table_name}.csv")
-        cursor.execute(f"SELECT * FROM {table_name};")
+        # check if table_name is not malicious
+        if not is_safe_table_name(table_name):
+            print('Malicious table name detected')
+            return 1
+        csv_file = os.path.join(os.path.dirname(os.path.realpath(__file__)), f'{table_name}.csv')
+        cursor.execute('SELECT * FROM %s;' % table_name)
         rows = cursor.fetchall()
         with open(csv_file, 'w', newline='\n') as file:
             writer = csv.writer(file, dialect='unix', escapechar='\\')
@@ -44,14 +53,18 @@ def export_tables_mariadb_to_csv(config):
     cursor = conn.cursor()
 
     # Get the list of tables in the database
-    cursor.execute("SHOW TABLES;")
+    cursor.execute('SHOW TABLES;')
     tables = cursor.fetchall()
 
     # Export each table to a separate CSV file
     for table in tables:
         table_name = table[0]
-        csv_file = os.path.join(os.path.dirname(os.path.realpath(__file__)), f"{table_name}.csv.mariadb")
-        cursor.execute(f"SELECT * FROM {table_name};")
+        # check if table_name is not malicious
+        if not is_safe_table_name(table_name):
+            print('Malicious table name detected')
+            return 1
+        csv_file = os.path.join(os.path.dirname(os.path.realpath(__file__)), f'{table_name}.csv.mariadb')
+        cursor.execute('SELECT * FROM %s;' % table_name)
         rows = cursor.fetchall()
 
         with open(csv_file, 'w', newline='\n') as file:

install.sh

@@ -7,7 +7,7 @@ LINUX_PACKAGE_MANAGER="apt-get"
 
 install_linux_packages() {
     # Install required packages
-    PACKAGES="python3 python3-pip wget curl sqlite3 libsqlite3-dev cmake gcc sudo apt-get install libmariadb-dev jq"
+    PACKAGES="python3 python3-pip wget curl sqlite3 libsqlite3-dev cmake gcc libmariadb-dev jq"
     which ${LINUX_PACKAGE_MANAGER} &> /dev/null
     if [ $? != 0 ]; then
         printf "${RED}Could not find ${LINUX_PACKAGE_MANAGER} command.\\nPlease specify your package manager at the start of the script.\\n${SANE}"
@@ -50,7 +50,7 @@ setup_create_db() {
     fi
     cd ".."
 
-    ## configure submodules of SQLiteCpp for create_db
+    ## configure submodules of mariadb-connector-cpp for create_db
     cd "mariadb-connector-cpp"
     if [ $QUIET != 1 ]; then
         git submodule init

migrate_sqlite_to_mariadb.sh

@@ -57,17 +57,17 @@ PASSWORD=$(jq -r '.DATABASE.PASSWORD' $CONFIG_FILE)
 PORT=$(jq -r '.DATABASE.PORT' $CONFIG_FILE)
 DATABASE_NAME=$(jq -r '.DATABASE_NAME' $CONFIG_FILE)
 CPE_DATABASE_NAME=$(jq -r '.cpe_search.DATABASE_NAME' $CONFIG_FILE)
-CREATE_TABLES_QUERIES_VULNDB=$ABS_PATH/$(jq -r '.CREATE_SQL_STATEMENTS_FILE' $CONFIG_FILE)
-CREATE_TABLES_QUERIES_CPE_SEARCH=$ABS_PATH/$(jq -r '.cpe_search.CREATE_SQL_STATEMENTS_FILE' $CONFIG_FILE)
+CREATE_TABLES_QUERIES_VULNDB=$ABS_PATH/create_sql_statements.json
+CREATE_TABLES_QUERIES_CPE_SEARCH=$ABS_PATH/cpe_search/create_sql_statements.json
 
 
 # Export sqlite databases
-echo "[+] Export sqlite as csv"
-python3 $ABS_PATH/export_database_as_csv.py sqlite $DATABASE_FILE
-python3 $ABS_PATH/export_database_as_csv.py sqlite $CPE_DATABASE_FILE
+echo "[+] Export SQLite as csv"
+python3 $ABS_PATH/export_database_as_csv.py sqlite $DATABASE_FILE || { rm $ABS_PATH/*.csv; echo "[-] Migration failed"; exit 1; }
+python3 $ABS_PATH/export_database_as_csv.py sqlite $CPE_DATABASE_FILE || { rm $ABS_PATH/*.csv; echo "[-] Migration failed"; exit 1; }
 
 # Create databases
-echo "[+] Add data to mariadb"
+echo "[+] Add data to MariaDB"
 # get queries from file
 vulndb_create_tables_queries=$(cat $CREATE_TABLES_QUERIES_VULNDB | jq '.TABLES | .[] | select(.mariadb) | .mariadb'| tr '\n' ' ' | sed 's/"//g')
 cpe_search_create_tables_queries=$(cat $CREATE_TABLES_QUERIES_CPE_SEARCH | jq '.TABLES | .[] | select(.mariadb) | .mariadb' | tr '\n' ' ' | sed 's/"//g')
@@ -85,9 +85,9 @@ vulndb_create_views=$(cat $CREATE_TABLES_QUERIES_VULNDB | jq '.VIEWS | .[] | sel
 mariadb -u $USER --password=$PASSWORD -h $HOST -P $PORT -D "$DATABASE_NAME" -e "$vulndb_create_views"
 
 # Export mariadb databases
-echo "[+] Export mariadb as csv"
-python3 $ABS_PATH/export_database_as_csv.py mariadb $DATABASE_NAME,$USER,$PASSWORD,$HOST,$PORT
-python3 $ABS_PATH/export_database_as_csv.py mariadb $CPE_DATABASE_NAME,$USER,$PASSWORD,$HOST,$PORT
+echo "[+] Export MariaDB as csv"
+python3 $ABS_PATH/export_database_as_csv.py mariadb $DATABASE_NAME,$USER,$PASSWORD,$HOST,$PORT || { rm $ABS_PATH/*.csv $ABS_PATH/*.csv.mariadb; echo "[-] Migration failed"; exit 1; }
+python3 $ABS_PATH/export_database_as_csv.py mariadb $CPE_DATABASE_NAME,$USER,$PASSWORD,$HOST,$PORT|| { rm $ABS_PATH/*.csv $ABS_PATH/*.csv.mariadb; echo "[-] Migration failed"; exit 1; }
 
 # check whether everything migrated correctly
 # Loop through each CSV file in the current folder

search_vulns.py

@@ -125,9 +125,9 @@ def get_vuln_details(db_cursor, vulns, add_other_exploit_refs):
         query = 'SELECT edb_ids, description, published, last_modified, cvss_version, base_score, vector FROM cve WHERE cve_id = ?'
         db_cursor.execute(query, (cve_id,))
         edb_ids, descr, publ, last_mod, cvss_ver, score, vector = db_cursor.fetchone()
-        detailed_vulns[cve_id] = {"id": cve_id, "description": descr, "published": publ, "modified": last_mod,
-                                  "href": "https://nvd.nist.gov/vuln/detail/%s" % cve_id, "cvss_ver": cvss_ver,
-                                  "cvss": score, "cvss_vec": vector, 'vuln_match_reason': match_reason}
+        detailed_vulns[cve_id] = {"id": cve_id, "description": descr, "published": str(publ), "modified": str(last_mod),
+                                  "href": "https://nvd.nist.gov/vuln/detail/%s" % cve_id, "cvss_ver": str(float(cvss_ver)),
+                                  "cvss": str(float(score)), "cvss_vec": vector, 'vuln_match_reason': match_reason}
 
         edb_ids = edb_ids.strip()
         if edb_ids:
@@ -290,7 +290,7 @@ def print_vulns(vulns, to_string=False):
                     print_str += len("Exploits:  ") * " " + edb_link + "\n"
 
         print_str += "Reference: " + vuln_node["href"]
-        print_str += ", " + str(vuln_node["published"]).split(" ")[0]
+        print_str += ", " + vuln_node["published"].split(" ")[0]
         if not to_string:
             printit(print_str)
         else:
@@ -660,4 +660,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()