From ba0d820a4590545582e2bf041a9963ba3040528d Mon Sep 17 00:00:00 2001
From: johackim <contact@johackim.com>
Date: Sat, 6 Apr 2024 02:18:27 +0200
Subject: [PATCH] feat: add whitelist time & password support

---
 src/commands.js | 16 +++++++++++++++-
 src/daemon.js   |  8 +++++---
 src/utils.js    | 10 ++++++++++
 3 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/src/commands.js b/src/commands.js
index f5d7eb4..318253f 100644
--- a/src/commands.js
+++ b/src/commands.js
@@ -60,6 +60,20 @@ export const unblockCmd = (name) => {
 };
 
 export const whitelistCmd = (name) => {
+    const time = getParam('--time') || getParam('-t');
+    const password = getParam('--password') || getParam('-p');
+    const distraction = { name, time };
+
+    if (isValidPassword(password)) {
+        disableShieldMode(password);
+        console.log('Shield mode disabled.');
+        whitelistDistraction(distraction);
+        console.log(`Whitelisting ${name}`);
+        enableShieldMode(password);
+        console.log('Shield mode enabled.');
+        return;
+    }
+
     if (config?.shield) {
         console.log('You must disable the shield mode first.');
         return;
@@ -70,7 +84,7 @@ export const whitelistCmd = (name) => {
         return;
     }
 
-    whitelistDistraction({ name });
+    whitelistDistraction(distraction);
     console.log(`Whitelisting ${name}`);
 };
 
diff --git a/src/daemon.js b/src/daemon.js
index 8f88d6d..01b7abc 100644
--- a/src/daemon.js
+++ b/src/daemon.js
@@ -24,14 +24,16 @@ const handleAppBlocking = () => {
 };
 
 const handleTimeout = () => {
-    const blocklist = config?.blocklist.filter(({ timeout }) => {
+    const list = [...config.blocklist, ...config.whitelist];
+
+    const listWithoutTimeout = list.filter(({ timeout }) => {
         if (!timeout) return true;
         return timeout >= Math.floor(Date.now() / 1000);
     });
 
-    if (blocklist?.length === config?.blocklist.length) return;
+    if (listWithoutTimeout.length === list.length) return;
 
-    editConfig({ ...config, blocklist });
+    editConfig(config);
 };
 
 const cleanUpAndExit = () => {
diff --git a/src/utils.js b/src/utils.js
index 2a7b67c..b254995 100644
--- a/src/utils.js
+++ b/src/utils.js
@@ -97,6 +97,7 @@ export const editConfig = ({ blocklist = [], whitelist = [], shield, password, p
     config.whitelist = removeDuplicates(config.shield ? config.whitelist : whitelist);
     config.blocklist = removeDuplicates(config.shield ? [...config.blocklist, ...blocklist] : blocklist);
     config.blocklist = config.blocklist.filter(({ timeout }) => !timeout || timeout >= Math.floor(Date.now() / 1000));
+    config.whitelist = config.whitelist.filter(({ timeout }) => !timeout || timeout >= Math.floor(Date.now() / 1000));
 
     if (isValidPassword(password)) {
         unblockRoot();
@@ -191,6 +192,13 @@ export const unblockDistraction = (distraction) => {
 
 export const whitelistDistraction = (distraction) => {
     config.whitelist.push(distraction);
+    config.whitelist = config.whitelist.map((d) => {
+        if (getTimeType(d.time) === 'duration') {
+            return { ...d, timeout: createTimeout(d.time) };
+        }
+
+        return d;
+    });
     sendDataToSocket(config);
 };
 
@@ -198,6 +206,7 @@ export const getRootDomain = (domain) => domain.split('.').slice(-2).join('.');
 
 export const isDistractionWhitelisted = (distraction) => {
     if (config.whitelist.some((d) => d.name === distraction)) return true;
+    if (config.whitelist.some((d) => d.name === '*')) return true;
     if (config.whitelist.some((d) => d.name === `*.${getRootDomain(distraction)}`)) return true;
 
     return false;
@@ -241,6 +250,7 @@ export const sendNotification = (title, message) => {
 export const getRunningBlockedApps = () => {
     const blockedApps = config.blocklist
         .filter(({ name }) => !isValidDomain(name) && !name.includes('*'))
+        .filter(({ name }) => !isDistractionWhitelisted(name))
         .filter(({ time }) => isWithinTimeRange(time))
         .map(({ name }) => name);