From 4f4a12bb40ab6f09a8651ad5ad908edc8c9b59b1 Mon Sep 17 00:00:00 2001
From: DL6ER <dl6er@dl6er.de>
Date: Fri, 5 Jul 2019 14:03:57 +0200
Subject: [PATCH] Upgrade database if necessary and store audit domains
 therein.

Signed-off-by: DL6ER <dl6er@dl6er.de>
---
 .../Scripts/database_migration/gravity-db.sh  |  4 +-
 gravity.sh                                    | 54 +++++++++++++------
 2 files changed, 40 insertions(+), 18 deletions(-)

diff --git a/advanced/Scripts/database_migration/gravity-db.sh b/advanced/Scripts/database_migration/gravity-db.sh
index 2c5669f0..492546be 100644
--- a/advanced/Scripts/database_migration/gravity-db.sh
+++ b/advanced/Scripts/database_migration/gravity-db.sh
@@ -11,8 +11,8 @@
 # Please see LICENSE file for your rights under this license.
 
 upgrade_gravityDB(){
-	version=$(sqlite3 "$1" "SELECT "value" FROM "info" WHERE "property" = 'version';")
-	echo $version
+	local version=$(sqlite3 "$1" "SELECT "value" FROM "info" WHERE "property" = 'version';")
+
 	case "$version" in
 	1)
 		sqlite3 "$1" < "/etc/.pihole/advanced/Scripts/database_migration/gravity/1_to_2.sql"
diff --git a/gravity.sh b/gravity.sh
index 1a325ba7..724238fa 100755
--- a/gravity.sh
+++ b/gravity.sh
@@ -17,6 +17,8 @@ coltable="/opt/pihole/COL_TABLE"
 source "${coltable}"
 regexconverter="/opt/pihole/wildcard_regex_converter.sh"
 source "${regexconverter}"
+readonly databaseMigrationScript="/etc/.pihole/advanced/Scripts/database_migration/gravity-db.sh"
+source "${databaseMigrationScript}"
 
 basename="pihole"
 PIHOLE_COMMAND="/usr/local/bin/${basename}"
@@ -28,6 +30,7 @@ whitelistFile="${piholeDir}/whitelist.txt"
 blacklistFile="${piholeDir}/blacklist.txt"
 regexFile="${piholeDir}/regex.list"
 adListFile="${piholeDir}/adlists.list"
+auditFile="${piholeDir}/audit.list"
 
 localList="${piholeDir}/local.list"
 VPNList="/etc/openvpn/ipp.txt"
@@ -116,14 +119,27 @@ database_table_from_file() {
     declare -i rowid
     rowid=1
     # Read file line by line
-    grep -v '^ *#' < "${source}" | while IFS= read -r domain
-    do
-      # Only add non-empty lines
-      if [[ ! -z "${domain}" ]]; then
-        echo "${rowid},\"${domain}\",1,${timestamp},${timestamp},\"Migrated from ${source}\"" >> "${tmpFile}"
-        rowid+=1
-      fi
-    done
+    if [[ "${table}" == "audit" ]]; then
+      grep -v '^ *#' < "${source}" | while IFS= read -r domain
+      do
+        # Only add non-empty lines
+        if [[ ! -z "${domain}" ]]; then
+          # Audit table format
+          echo "${rowid},\"${domain}\",${timestamp}" >> "${tmpFile}"
+          rowid+=1
+        fi
+      done
+    else
+      grep -v '^ *#' < "${source}" | while IFS= read -r domain
+      do
+        # Only add non-empty lines
+        if [[ ! -z "${domain}" ]]; then
+          # White-, black-, and regexlist format
+          echo "${rowid},\"${domain}\",1,${timestamp},${timestamp},\"Migrated from ${source}\"" >> "${tmpFile}"
+          rowid+=1
+        fi
+      done
+    fi
     inputfile="${tmpFile}"
   fi
   # Store domains in database table specified by ${table}
@@ -150,34 +166,40 @@ database_table_from_file() {
 # Migrate pre-v5.0 list files to database-based Pi-hole versions
 migrate_to_database() {
   # Create database file only if not present
-  if [ -e "${gravityDBfile}" ]; then
-    return 0
+  if [ ! -e "${gravityDBfile}" ]; then
+    echo -e "  ${INFO} Creating new gravity database"
+    generate_gravity_database
   fi
 
-  echo -e "  ${INFO} Creating new gravity database"
-  generate_gravity_database
+  # Check if gravity database needs to be updated
+  upgrade_gravityDB "${gravityDBfile}"
 
   # Migrate list files to new database
-  if [[ -e "${adListFile}" ]]; then
+  if [ -e "${adListFile}" ]; then
     # Store adlist domains in database
     echo -e "  ${INFO} Migrating content of ${adListFile} into new database"
     database_table_from_file "adlist" "${adListFile}"
   fi
-  if [[ -e "${blacklistFile}" ]]; then
+  if [ -e "${blacklistFile}" ]; then
     # Store blacklisted domains in database
     echo -e "  ${INFO} Migrating content of ${blacklistFile} into new database"
     database_table_from_file "blacklist" "${blacklistFile}"
   fi
-  if [[ -e "${whitelistFile}" ]]; then
+  if [ -e "${whitelistFile}" ]; then
     # Store whitelisted domains in database
     echo -e "  ${INFO} Migrating content of ${whitelistFile} into new database"
     database_table_from_file "whitelist" "${whitelistFile}"
   fi
-  if [[ -e "${regexFile}" ]]; then
+  if [ -e "${regexFile}" ]; then
     # Store regex domains in database
     echo -e "  ${INFO} Migrating content of ${regexFile} into new database"
     database_table_from_file "regex" "${regexFile}"
   fi
+  if [ -e "${auditFile}" ]; then
+    # Store audit domains in database
+    echo -e "  ${INFO} Migrating content of ${auditFile} into new database"
+    database_table_from_file "audit" "${auditFile}"
+  fi
 }
 
 # Determine if DNS resolution is available before proceeding