From 190ab796063a53ac9843fd97cc8f978d10bd4a56 Mon Sep 17 00:00:00 2001
From: DL6ER <dl6er@dl6er.de>
Date: Fri, 20 Aug 2021 20:48:57 +0200
Subject: [PATCH] Implement fully-automated gravity database recovery method.

Signed-off-by: DL6ER <dl6er@dl6er.de>
---
 gravity.sh | 64 ++++++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 62 insertions(+), 2 deletions(-)

diff --git a/gravity.sh b/gravity.sh
index b2cd27b9..911242f7 100755
--- a/gravity.sh
+++ b/gravity.sh
@@ -848,6 +848,42 @@ gravity_Cleanup() {
   fi
 }
 
+database_recovery() {
+  str="Checking integrity of existing gravity database"
+  echo -ne "  ${INFO} ${str}..."
+  if result="$(pihole-FTL sqlite3 "${gravityDBfile}" "PRAGMA integrity_check" 2>&1)"; then
+    echo -e "${OVER}  ${TICK} ${str} - no errors found"
+  else
+    echo -e "${OVER}  ${CROSS} ${str} - errors found:"
+    while IFS= read -r line ; do echo "  - $line"; done <<< "$result"
+  fi
+
+  str="Checking foreign keys of existing gravity database"
+  echo -ne "  ${INFO} ${str}..."
+  if result="$(pihole-FTL sqlite3 "${gravityDBfile}" "PRAGMA foreign_key_check" 2>&1)"; then
+    echo -e "${OVER}  ${TICK} ${str} - no errors found"
+  else
+    echo -e "${OVER}  ${CROSS} ${str} - errors found:"
+    while IFS= read -r line ; do echo "  - $line"; done <<< "$result"
+  fi
+
+  str="Trying to recover existing gravity database"
+  echo -ne "  ${INFO} ${str}..."
+  # We have to remove any possibly existing recovery database or this will fail
+  rm -f "${gravityDBfile}.recovered" > /dev/null 2>&1
+  if result="$(pihole-FTL sqlite3 "${gravityDBfile}" ".recover" | pihole-FTL sqlite3 "${gravityDBfile}.recovered" 2>&1)"; then
+    echo -e "${OVER}  ${TICK} ${str} - success"
+    mv "${gravityDBfile}" "${gravityDBfile}.old"
+    mv "${gravityDBfile}.recovered" "${gravityDBfile}"
+  else
+    echo -e "${OVER}  ${CROSS} ${str} - the following errors happened:"
+    while IFS= read -r line ; do echo "  - $line"; done <<< "$result"
+    echo -e "  ${CROSS} Recovery failed. Try \"pihole -r recreate\" instead."
+    exit 1
+  fi
+  echo ""
+}
+
 helpFunc() {
   echo "Usage: pihole -g
 Update domains from blocklists specified in adlists.list
@@ -858,10 +894,30 @@ Options:
   exit 0
 }
 
+repairSelector() {
+  case "$1" in
+    "recover") recover_database=true;;
+    "recreate") recreate_database=true;;
+    *) echo "Usage: pihole -g -r {recover,recreate}
+Attempt to repair gravity database
+
+Available options:
+  pihole -g -r recover    Try to recover a damaged gravity database file.
+                          Pi-hole tries to restore as much as possible
+                          from a corrupted gravity database.
+  pihole -g -r recreate   Create a new gravity database file from scratch.
+                          This will remove your existing gravity database
+                          and create a new file from scratch. If you still
+                          have the migration backup created when migrating
+                          to Pi-hole v5.0, Pi-hole will import these files."
+    exit 0;;
+  esac
+}
+
 for var in "$@"; do
   case "${var}" in
     "-f" | "--force" ) forceDelete=true;;
-    "-r" | "--recreate" ) recreate_database=true;;
+    "-r" | "--repair" ) repairSelector "$3";;
     "-h" | "--help" ) helpFunc;;
   esac
 done
@@ -875,7 +931,7 @@ fi
 gravity_Trap
 
 if [[ "${recreate_database:-}" == true ]]; then
-  str="Restoring from migration backup"
+  str="Recreating gravity database from migration backup"
   echo -ne "${INFO} ${str}..."
   rm "${gravityDBfile}"
   pushd "${piholeDir}" > /dev/null || exit
@@ -884,6 +940,10 @@ if [[ "${recreate_database:-}" == true ]]; then
   echo -e "${OVER}  ${TICK} ${str}"
 fi
 
+if [[ "${recover_database:-}" == true ]]; then
+  database_recovery
+fi
+
 # Move possibly existing legacy files to the gravity database
 migrate_to_database