From 09c4c88a6d1794a3ddb3ee8acb96e2cfb03ad7c2 Mon Sep 17 00:00:00 2001
From: DL6ER <dl6er@dl6er.de>
Date: Sun, 3 Feb 2019 13:04:31 +0100
Subject: [PATCH] Create and fill gravity.db during pihole -g

Signed-off-by: DL6ER <dl6er@dl6er.de>
---
 advanced/Templates/gravity.db.schema | 13 +++++++
 gravity.sh                           | 56 ++++++++++++++++++++++++++--
 2 files changed, 65 insertions(+), 4 deletions(-)
 create mode 100644 advanced/Templates/gravity.db.schema

diff --git a/advanced/Templates/gravity.db.schema b/advanced/Templates/gravity.db.schema
new file mode 100644
index 00000000..1eee8491
--- /dev/null
+++ b/advanced/Templates/gravity.db.schema
@@ -0,0 +1,13 @@
+CREATE TABLE whitelist (domain TEXT UNIQUE NOT NULL, comment TEXT, DateAdded DATETIME);
+
+CREATE TABLE blacklist (domain TEXT UNIQUE NOT NULL, comment TEXT, DateAdded DATETIME);
+
+CREATE TABLE gravity   (domain TEXT UNIQUE NOT NULL);
+
+CREATE VIEW vw_gravity AS SELECT DISTINCT a.domain
+FROM gravity a
+WHERE a.domain NOT IN (SELECT domain from whitelist);
+
+CREATE VIEW vw_blacklist AS SELECT DISTINCT a.domain
+FROM blacklist a
+WHERE a.domain NOT IN (SELECT domain from whitelist);
diff --git a/gravity.sh b/gravity.sh
index 75a51608..ac045a93 100755
--- a/gravity.sh
+++ b/gravity.sh
@@ -35,6 +35,10 @@ blackList="${piholeDir}/black.list"
 localList="${piholeDir}/local.list"
 VPNList="/etc/openvpn/ipp.txt"
 
+piholeGitDir="/etc/.pihole"
+gravityDBfile="${piholeDir}/gravity.db"
+gravityDBschema="${piholeGitDir}/advanced/Templates/gravity.db.schema"
+
 domainsExtension="domains"
 matterAndLight="${basename}.0.matterandlight.txt"
 parsedMatter="${basename}.1.parsedmatter.txt"
@@ -83,6 +87,11 @@ if [[ -r "${piholeDir}/pihole.conf" ]]; then
   echo -e "  ${COL_LIGHT_RED}Ignoring overrides specified within pihole.conf! ${COL_NC}"
 fi
 
+# Generate new sqlite3 file from schema template
+generate_gravity_database() {
+  sqlite3 "${gravityDBfile}" < "${gravityDBschema}"
+}
+
 # Determine if Pi-hole blocking is disabled
 # If this is the case, we want to update
 #  gravity.list.bck and black.list.bck instead of
@@ -582,12 +591,38 @@ gravity_ParseBlacklistDomains() {
     cp "${piholeDir}/${preEventHorizon}" "${piholeDir}/${accretionDisc}"
   fi
 
-  # Move the file over as /etc/pihole/gravity.list so dnsmasq can use it
-  output=$( { mv "${piholeDir}/${accretionDisc}" "${adList}"; } 2>&1 )
+  # Create database file if not present
+  if [ ! -e "${gravityDBfile}" ]; then
+    generate_gravity_database
+  fi
+
+  # Backup gravity database
+  cp "${gravityDBfile}" "${gravityDBfile}.bck"
+
+  # Empty domains
+  output=$( { sqlite3 "${gravityDBfile}" <<< "DELETE FROM gravity;"; } 2>&1 )
   status="$?"
 
   if [[ "${status}" -ne 0 ]]; then
-    echo -e "\\n  ${CROSS} Unable to move ${accretionDisc} from ${piholeDir}\\n  ${output}"
+    echo -e "\\n  ${CROSS} Unable to truncate gravity database ${gravityDBfile}\\n  ${output}"
+    gravity_Cleanup "error"
+  fi
+
+  # Store domains in gravity database
+  output=$( { sqlite3 "${gravityDBfile}" <<< ".import ${piholeDir}/${accretionDisc} gravity"; } 2>&1 )
+  status="$?"
+
+  if [[ "${status}" -ne 0 ]]; then
+    echo -e "\\n  ${CROSS} Unable to create gravity database ${gravityDBfile}\\n  ${output}"
+    gravity_Cleanup "error"
+  fi
+
+  # Empty $adList if it already exists, otherwise, create it
+  output=$( { : > "${adList}"; } 2>&1 )
+  status="$?"
+
+  if [[ "${status}" -ne 0 ]]; then
+    echo -e "\\n  ${CROSS} Unable to create empty ${adList}\\n  ${output}"
     gravity_Cleanup "error"
   fi
 }
@@ -633,6 +668,19 @@ gravity_Cleanup() {
 
   echo -e "${OVER}  ${TICK} ${str}"
 
+  str="Optimizing domains database"
+  echo -ne "  ${INFO} ${str}..."
+  # Store
+  output=$( { sqlite3 "${gravityDBfile}" <<< "VACUUM;"; } 2>&1 )
+  status="$?"
+
+  if [[ "${status}" -ne 0 ]]; then
+    echo -e "\\n  ${CROSS} Unable to optimize gravity database ${gravityDBfile}\\n  ${output}"
+    gravity_Cleanup "error"
+  else
+    echo -e "${OVER}  ${TICK} ${str}"
+  fi
+
   # Only restart DNS service if offline
   if ! pidof ${resolver} &> /dev/null; then
     "${PIHOLE_COMMAND}" restartdns
@@ -707,7 +755,7 @@ gravity_ShowBlockCount
 
 # Perform when downloading blocklists, or modifying the white/blacklist (not wildcards)
 if [[ "${skipDownload}" == false ]] || [[ "${listType}" == *"list" ]]; then
-  str="Parsing domains into hosts format"
+  str="Parsing domains"
   echo -ne "  ${INFO} ${str}..."
 
   gravity_ParseUserDomains