From 037d52104a84ac867842e017c30879d17ee79cfe Mon Sep 17 00:00:00 2001
From: DL6ER <dl6er@dl6er.de>
Date: Tue, 26 Nov 2019 10:58:39 +0100
Subject: [PATCH] New command "pihole -g -r" recreates gravity.db based on
 files backed up in /etc/pihole/migration_update. This is useful to restore a
 working version of the database when the user destroyed the original
 database. Also, update gravity.db to version 5 because of a fix we needed to
 implement.

Signed-off-by: DL6ER <dl6er@dl6er.de>
---
 .../Scripts/database_migration/gravity-db.sh    | 12 ++++++++++--
 .../database_migration/gravity/4_to_5.sql       | 17 +++++++++++++++++
 gravity.sh                                      | 17 ++++++++++++++---
 3 files changed, 41 insertions(+), 5 deletions(-)
 create mode 100644 advanced/Scripts/database_migration/gravity/4_to_5.sql

diff --git a/advanced/Scripts/database_migration/gravity-db.sh b/advanced/Scripts/database_migration/gravity-db.sh
index 773898b0..1fe6a61f 100644
--- a/advanced/Scripts/database_migration/gravity-db.sh
+++ b/advanced/Scripts/database_migration/gravity-db.sh
@@ -44,9 +44,17 @@ upgrade_gravityDB(){
 		version=3
 	fi
 	if [[ "$version" == "3" ]]; then
-		# This migration script upgrades the gravity and adlist views
+		# This migration script upgrades the gravity and list views
 		# implementing necessary changes for per-client blocking
+		echo -e "  ${INFO} Upgrading gravity database from version 3 to 4"
 		sqlite3 "${database}" < "${scriptPath}/3_to_4.sql"
-		version=3
+		version=4
+	fi
+	if [[ "$version" == "4" ]]; then
+		# This migration script upgrades the adlist view
+		# to return an ID used in gravity.sh
+		echo -e "  ${INFO} Upgrading gravity database from version 4 to 5"
+		sqlite3 "${database}" < "${scriptPath}/4_to_5.sql"
+		version=5
 	fi
 }
diff --git a/advanced/Scripts/database_migration/gravity/4_to_5.sql b/advanced/Scripts/database_migration/gravity/4_to_5.sql
new file mode 100644
index 00000000..22b75d58
--- /dev/null
+++ b/advanced/Scripts/database_migration/gravity/4_to_5.sql
@@ -0,0 +1,17 @@
+.timeout 30000
+
+PRAGMA FOREIGN_KEYS=OFF;
+
+BEGIN TRANSACTION;
+
+DROP VIEW vw_adlist;
+CREATE VIEW vw_adlist AS SELECT DISTINCT address, adlist.id AS id
+    FROM adlist
+    LEFT JOIN adlist_by_group ON adlist_by_group.adlist_id = adlist.id
+    LEFT JOIN "group" ON "group".id = adlist_by_group.group_id
+    WHERE adlist.enabled = 1 AND (adlist_by_group.group_id IS NULL OR "group".enabled = 1)
+    ORDER BY adlist.id;
+
+UPDATE info SET value = 5 WHERE property = 'version';
+
+COMMIT;
diff --git a/gravity.sh b/gravity.sh
index 7a9e4f67..3225b28d 100755
--- a/gravity.sh
+++ b/gravity.sh
@@ -229,7 +229,7 @@ gravity_CheckDNSResolutionAvailable() {
   fi
 
   # Determine if $lookupDomain is resolvable
-  if timeout 1 getent hosts "${lookupDomain}" &> /dev/null; then
+  if timeout 4 getent hosts "${lookupDomain}" &> /dev/null; then
     # Print confirmation of resolvability if it had previously failed
     if [[ -n "${secs:-}" ]]; then
       echo -e "${OVER}  ${TICK} DNS resolution is now available\\n"
@@ -243,7 +243,7 @@ gravity_CheckDNSResolutionAvailable() {
   # If the /etc/resolv.conf contains resolvers other than 127.0.0.1 then the local dnsmasq will not be queried and pi.hole is NXDOMAIN.
   # This means that even though name resolution is working, the getent hosts check fails and the holddown timer keeps ticking and eventualy fails
   # So we check the output of the last command and if it failed, attempt to use dig +short as a fallback
-  if timeout 1 dig +short "${lookupDomain}" &> /dev/null; then
+  if timeout 4 dig +short "${lookupDomain}" &> /dev/null; then
     if [[ -n "${secs:-}" ]]; then
       echo -e "${OVER}  ${TICK} DNS resolution is now available\\n"
     fi
@@ -425,7 +425,7 @@ gravity_DownloadBlocklistFromUrl() {
   if [[ "${success}" == true ]]; then
     if [[ "${httpCode}" == "304" ]]; then
       # Add domains to database table
-      str="Adding to database table"
+      str="Adding adlist with ID ${adlistID} to database table"
       echo -ne "  ${INFO} ${str}..."
       database_table_from_file "gravity" "${saveLocation}" "${adlistID}"
       echo -e "${OVER}  ${TICK} ${str}"
@@ -660,6 +660,7 @@ for var in "$@"; do
   case "${var}" in
     "-f" | "--force" ) forceDelete=true;;
     "-o" | "--optimize" ) optimize_database=true;;
+    "-r" | "--recreate" ) recreate_database=true;;
     "-h" | "--help" ) helpFunc;;
   esac
 done
@@ -667,6 +668,16 @@ done
 # Trap Ctrl-C
 gravity_Trap
 
+if [[ "${recreate_database:-}" == true ]]; then
+  str="Restoring from migration backup"
+  echo -ne "${INFO} ${str}..."
+  rm "${gravityDBfile}"
+  pushd "${piholeDir}" > /dev/null
+  cp migration_backup/* .
+  popd > /dev/null
+  echo -e "${OVER}  ${TICK} ${str}"
+fi
+
 # Move possibly existing legacy files to the gravity database
 migrate_to_database