Refactor improve & lose funcs; work on takeBreak

main.go

@@ -4,22 +4,21 @@ import (
 	"bufio"
 	"errors"
 	"fmt"
-	"io/ioutil"
+	"io"
 	"log"
 	"os"
 	"os/signal"
 	"os/user"
 	"runtime"
-	"strings"
 	"time"
 
 	"github.com/akamensky/argparse"
 )
 
 var (
-	marker                string = "## MANAGED BY HYPERFOCUS ##\n"
 	distractorsFile       string = "/etc/hyperfocus_distractors"
 	predefDistractorsFile string = "/etc/hyperfocus_distractors_predefined"
+	hostBackupFile        string = "/etc/.backup_hosts"
 )
 
 const resetDNSMsg string = `
@@ -30,12 +29,35 @@ macOS: open terminal and run "sudo dscacheutil -flushcache"
 Linux with Systemd: open terminal and run "sudo systemctl restart systemd-resolved"
 Other init systems or *BSD: You probably already know how to do this.`
 
+// Check if <path> exists on the host filesystem
 func fileExists(path string) bool {
 
 	_, err := os.Stat(path)
 	return !errors.Is(err, os.ErrNotExist)
 }
 
+// Copy src to dest
+func copyFile(src string, dest string) error {
+
+	origFile, err := os.Open(src)
+	if err != nil {
+		return fmt.Errorf("could not open original file %s for copying: %s", src, err)
+	}
+
+	backupFile, err := os.OpenFile(dest, os.O_CREATE|os.O_WRONLY, 0666)
+	if err != nil {
+		return fmt.Errorf("could not open backup file %s for writing: %s", dest, err)
+	}
+
+	_, err = io.Copy(backupFile, origFile)
+	if err != nil {
+		return fmt.Errorf("could not copy %s to %s: %s", src, dest, err)
+	}
+
+	return nil
+}
+
+// Return system's host file as string
 func getHostFile() string {
 
 	var hostFile string
@@ -50,6 +72,7 @@ func getHostFile() string {
 	return hostFile
 }
 
+// Check if running as super user
 func isRoot() bool {
 
 	currUser, err := user.Current()
@@ -63,21 +86,29 @@ func isRoot() bool {
 // Disables access to websites that are defined as 'distractors'
 func improveFocus() {
 
+	// Backup host file if a backup does not already exist
+	if !fileExists(hostBackupFile) {
+		if err := copyFile(getHostFile(), hostBackupFile); err != nil {
+			log.Fatalln(err)
+		}
+	}
+
+	// Open host file for writing/appending
 	hostFile, err := os.OpenFile(getHostFile(), os.O_APPEND|os.O_WRONLY, 0666)
 	if err != nil {
 		log.Fatalln(err)
 	}
 	defer hostFile.Close()
 
-	if _, err := hostFile.WriteString(marker); err != nil {
+	// For checking if /etc/hyperfocus_distractors or /etc/hyperfocus_distractors_predefined
-		log.Fatalln(err)
+	// exist.
-	}
-
 	var (
 		distractorsFileWarn       bool
 		predefDistractorsFileWarn bool
 	)
 
+	// If /etc/hyperfocus_distractors exists, take each host from it and append it
+	// to /etc/hosts for blocking. Else set distractorsFileWarn to true.
 	if fileExists(distractorsFile) {
 
 		distractorsFileObj, err := os.Open(distractorsFile)
@@ -97,6 +128,8 @@ func improveFocus() {
 		distractorsFileWarn = true
 	}
 
+	// If /etc/hyperfocus_distractors_predefined exists, take each host from it and,
+	// append it to /etc/hosts for blocking. Else set predefDistractorsFileWarn to true.
 	if fileExists(predefDistractorsFile) {
 
 		predefDistractorsFileObj, err := os.Open(predefDistractorsFile)
@@ -116,31 +149,25 @@ func improveFocus() {
 		predefDistractorsFileWarn = true
 	}
 
+	// Fail with warning if neither distractors files exist.
 	if distractorsFileWarn && predefDistractorsFileWarn {
-		log.Fatalln("Error: Please define a set of distractors in your distractors file, one per line. Alternatively, you can use a predefined set by running `hf predefined`.")
+		log.Fatalln("Error: Please define a set of distractors in your distractors file, one per line.",
-	}
+			"Alternatively, you can use a predefined set by running `sudo hf predefined`.")
-
-	if _, err := hostFile.WriteString(marker); err != nil {
-		log.Fatalln(err)
 	}
 
 	fmt.Println("Focus is now improved 😊")
-	fmt.Println(resetDNSMsg)
 }
 
 // Enables access to websites that are defined as 'distractors'
 func loseFocus() {
 
-	hostFileContent, err := ioutil.ReadFile(getHostFile())
+	// Remove the current /etc/hosts file before restoring the backup.
-	if err != nil {
+	if err := os.Remove(getHostFile()); err != nil {
 		log.Fatalln(err)
 	}
 
-	hostFileSlice := strings.Split(string(hostFileContent), marker)
+	// Restore the backup of /etc/hosts
-	hostFileSlice = append(hostFileSlice[:1], hostFileSlice[2:]...)
+	if err := copyFile(hostBackupFile, getHostFile()); err != nil {
-	hostFileBytes := []byte(strings.Join(hostFileSlice, ""))
-
-	if err := ioutil.WriteFile(getHostFile(), hostFileBytes, 0666); err != nil {
 		log.Fatalln(err)
 	}
 
@@ -154,19 +181,36 @@ func takeBreak(minutes int) {
 	fmt.Println("Your (probably) well-deserved break is commencing...")
 	loseFocus()
 
+	// Create a channel to receive os.Interrupt signal from user
 	c := make(chan os.Signal, 1)
 	signal.Notify(c, os.Interrupt)
 
-	sig := <-c
+	// TODO: Read up on Goroutines and channels and figure how to 'pipe' the os.Signal
-	t := minutes * 60
+	// channel to the goroutine below.
-	for i := t; i > 0; i-- {
-		fmt.Printf("%d seconds remaining", i)
-		time.Sleep(time.Second)
 
-		if sig == os.Interrupt {
+	// Prepare channel to receive signal. Sleep for t minus <minutes * 60> seconds,
-			break
+	// or break the countdown when interrupt signal is received from user, then
+	// improve focus again.
+	ticker := time.NewTicker(1 * time.Second)
+	done := make(chan bool)
+
+	go func() {
+		for {
+			select {
+			case <-done:
+				return
+			case t := <-ticker.C:
+				// TODO: Figure how to get the number of seconds as integers from t
+				fmt.Println(t)
+			case <-c:
+				return
 			}
 		}
+	}()
+
+	time.Sleep(time.Duration(minutes*60) * time.Second)
+	ticker.Stop()
+	done <- true
 
 	improveFocus()
 	fmt.Println("Welcome back to hyperfocus mode 😊 ")
@@ -175,9 +219,11 @@ func takeBreak(minutes int) {
 // Prints the current list of distractors to be blocked
 func listDistractors() {
 
-	userDistractorsFile, err := os.Open(distractorsFile)
+	for _, v := range []string{distractorsFile, predefDistractorsFile} {
+
+		userDistractorsFile, err := os.Open(v)
 		if err != nil {
-		log.Fatalln(err)
+			fmt.Println(err)
 		}
 		defer userDistractorsFile.Close()
 
@@ -188,6 +234,8 @@ func listDistractors() {
 		}
 	}
 
+}
+
 // Adds predefined distractors to .hyperfocus_distractors file
 func addPredefinedDistractors() {
 
@@ -240,6 +288,33 @@ func addPredefinedDistractors() {
 		"myparentsjoinedfacebook.com",
 		"myspace.com",
 		"netflix.com",
+		"news.anandtech.com",
+		"news.arstechnica.com",
+		"news.buzzfeed.com",
+		"news.cheezburger.com",
+		"news.cnet.com",
+		"news.cracked.com",
+		"news.distrowatch.com",
+		"news.ebay.com",
+		"news.facebook.com",
+		"news.gizmodo.com",
+		"news.homestarrunner.com",
+		"news.imgur.com",
+		"news.lifehacker.com",
+		"news.meetup.com",
+		"news.myspace.com",
+		"news.reddit.com",
+		"news.sciencemag.org",
+		"news.slashdot.com",
+		"news.slashdot.org",
+		"news.stripgenerator.com",
+		"news.theverge.com",
+		"news.trulia.com",
+		"news.typepad.com",
+		"news.ycombinator.com",
+		"news.ycombinator.com",
+		"news.yelp.com",
+		"news.youtube.com",
 		"notalwaysright.com",
 		"omegle.com",
 		"opentable.com",
@@ -279,7 +354,6 @@ func addPredefinedDistractors() {
 		"wired.com",
 		"woot.com",
 		"xkcd.com",
-		"ycombinator.com",
 		"yelp.com",
 		"youtube.com",
 		"zdnet.com",