hyperfocus/main.go

package main

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"os/signal"
	"os/user"
	"path/filepath"
	"runtime"
	"time"

	"github.com/akamensky/argparse"
)

var localFilePaths = make(map[string]string, 4)

// Check if <path> exists on the host filesystem
func fileExists(path string) bool {

	_, err := os.Stat(path)
	return !errors.Is(err, os.ErrNotExist)
}

func getLocalFilePaths() {

	switch runtime.GOOS {
	case "windows":
		localFilePaths["hosts"] = filepath.Clean("C:\\Windows\\System32\\drivers\\etc\\hosts")
		localFilePaths["hostsBackup"] = filepath.Clean("C:\\Windows\\System32\\drivers\\etc\\hosts.backup")
		localFilePaths["uDistractors"] = filepath.Clean(os.Getenv("HOME") + "\\.hf_distractors")
		localFilePaths["pDistractors"] = filepath.Clean(os.Getenv("HOME") + "\\.hf_predef_distractors")

	default:
		localFilePaths["hosts"] = filepath.Clean("/etc/hosts")
		localFilePaths["hostsBackup"] = filepath.Clean("/etc/.hosts.backup")
		localFilePaths["uDistractors"] = filepath.Clean("/etc/hf_distractors")
		localFilePaths["pDistractors"] = filepath.Clean("/etc/hf_predef_distractors")
	}
}

// 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
}

// Check if running as super user
func isRoot() bool {

	var isSu bool
	switch runtime.GOOS {
	case "windows":
		_, err := os.Open("\\\\.\\PHYSICALDRIVE0")
		if err != nil {
			isSu = false
		} else {
			isSu = true
		}

	default:
		currUser, err := user.Current()
		if err != nil {
			log.Fatalf("[Super user check] Unable to get current user: %s\n", err)
		}
		isSu = currUser.Username == "root"
	}

	return isSu
}

// Disables access to websites that are defined as 'distractors'
func improveFocus() {

	var displayMsg bool = false

	// If first run, display resetDNSMsg
	if !fileExists(localFilePaths["hostsBackup"]) {
		displayMsg = true
	}

	// Backup host file if a backup does not already exist
	if !fileExists(localFilePaths["hostsBackup"]) {
		if err := copyFile(localFilePaths["hosts"], localFilePaths["hostsBackup"]); err != nil {
			log.Fatalln(err)
		}
	}

	// Open host file for writing/appending
	hostFileObj, err := os.OpenFile(localFilePaths["hosts"], os.O_APPEND|os.O_WRONLY, 0666)
	if err != nil {
		log.Fatalln(err)
	}
	defer func() {
		if err := hostFileObj.Close(); err != nil {
			log.Printf("Error closing file: %s\n", err)
		}
	}()

	// For checking if /etc/hf_distractors or /etc/hf_distractors_predefined
	// exist.
	var (
		distractorsFileWarn       bool
		predefDistractorsFileWarn bool
	)

	// If /etc/hf_distractors exists, take each host from it and append it
	// to /etc/hosts for blocking. Else set distractorsFileWarn to true.
	if fileExists(localFilePaths["uDistractors"]) {

		distractorsFileObj, err := os.Open(localFilePaths["uDistractors"])
		if err != nil {
			log.Fatalln(err)
		}
		defer func() {
			if err := distractorsFileObj.Close(); err != nil {
				log.Printf("Error closing file: %s\n", err)
			}
		}()

		scanner := bufio.NewScanner(distractorsFileObj)
		for scanner.Scan() {
			var hostLine string = fmt.Sprintf("127.0.0.1\t%s\n", scanner.Text())
			if _, err := hostFileObj.WriteString(hostLine); err != nil {
				log.Fatalln(err)
			}
		}
	} else {
		distractorsFileWarn = true
	}

	// If /etc/hf_predef_distractors exists, take each host from it and,
	// append it to /etc/hosts for blocking. Else set predefDistractorsFileWarn to true.
	if fileExists(localFilePaths["pDistractors"]) {

		predefDistractorsFileObj, err := os.Open(localFilePaths["pDistractors"])
		if err != nil {
			log.Fatalln(err)
		}
		defer func() {
			if err := predefDistractorsFileObj.Close(); err != nil {
				log.Printf("Error closing file: %s\n", err)
			}
		}()

		scanner := bufio.NewScanner(predefDistractorsFileObj)
		for scanner.Scan() {
			var hostLine string = fmt.Sprintf("127.0.0.1\t%s\n", scanner.Text())
			if _, err := hostFileObj.WriteString(hostLine); err != nil {
				log.Fatalln(err)
			}
		}
	} else {
		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 `sudo hf predefined`.")
	}

	fmt.Println("Focus is now improved 😊")
	if displayMsg {
		fmt.Println(resetDNSMsg)
	}
}

// Enables access to websites that are defined as 'distractors'
func loseFocus() {

	// Remove the current /etc/hosts file before restoring the backup.
	if err := os.Remove(localFilePaths["hosts"]); err != nil {
		log.Fatalln(err)
	}

	// Restore the backup of /etc/hosts
	if err := copyFile(localFilePaths["hostsBackup"], localFilePaths["hosts"]); err != nil {
		log.Fatalln(err)
	}

	fmt.Println("Focus is now lost 🤪")
}

// Enables temporarily breaking concentration
func takeBreak(minutes int) {

	fmt.Println()
	fmt.Println("Your (probably) well-deserved break is commencing...")
	loseFocus()

	// Create notify context for os.Interrupt signal
	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
	defer cancel()

	// Create context with timeout for duration of minutes argument
	ctx, cancel = context.WithTimeout(ctx, time.Duration(minutes*60)*time.Second)
	defer cancel()

	<-ctx.Done()

	improveFocus()
}

// Prints the current list of distractors to be blocked
// If neither /etc/hf_distractors nor /etc/hf_predef_distractors exist,
// print error message to console but do not exit.
func listDistractors() {

	// Open /etc/hf_distractors and store it as *os.File type
	userDistractorsFileObj, err := os.Open(localFilePaths["uDistractors"])
	if err != nil {
		fmt.Println(err)
	}
	defer func() {
		if err := userDistractorsFileObj.Close(); err != nil {
			fmt.Printf("Error closing file: %s\n", err)
		}
	}()

	// Initialize a new scanner, scan /etc/hf_distractors, and print to
	// stdout line by line.
	scanner := bufio.NewScanner(userDistractorsFileObj)
	for scanner.Scan() {
		fmt.Println(scanner.Text())
	}

	// Open /etc/hf_predef_distractors and store it as *os.File type
	predefDistractorsFileObj, err := os.Open(localFilePaths["pDistractors"])
	if err != nil {
		fmt.Println(err)
	}
	defer func() {
		if err := predefDistractorsFileObj.Close(); err != nil {
			fmt.Printf("Error closing file: %s\n", err)
		}
	}()

	// Initialize a new scanner, scan /etc/hf_predef_distractors, and print to
	// stdout line by line.
	scanner = bufio.NewScanner(predefDistractorsFileObj)
	for scanner.Scan() {
		fmt.Println(scanner.Text())
	}

}

// Adds predefined distractors to /etc/hf_predef_distractors file
func addPredefinedDistractors() {

	// Slice of distractor domains
	distractors := []string{
		"anandtech.com",
		"answerbag.com",
		"arstechnica.com",
		"awkwardfamilyphotos.com",
		"bbc.com",
		"befunky.com",
		"break.com",
		"buzzfeed.com",
		"cheezburger.com",
		"cnet.com",
		"cnn.com",
		"cracked.com",
		"digg.com",
		"digitaltrends.com",
		"distrowatch.com",
		"ebay.com",
		"engadget.com",
		"everypoet.com",
		"facebook.com",
		"failblog.com",
		"fark.com",
		"firstpost.com",
		"flyertalk.com",
		"fmylife.com",
		"fox.com",
		"freerice.com",
		"gawker.com",
		"geek.com",
		"gizmag.com",
		"gizmodo.com",
		"gsmarena.com ",
		"homestarrunner.com",
		"hulu.com",
		"imgur.com",
		"imore.com",
		"instagram.com",
		"iwastesomuchtime.com",
		"lifehacker.com",
		"liveplasma.com",
		"lunchtimers.com",
		"mashable.com",
		"mcsweeneys.net",
		"medium.com",
		"meetup.com",
		"mono-1.com",
		"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",
		"passiveaggressivenotes.com",
		"pcworld.com",
		"petsinclothes.com",
		"phoronix.com",
		"pinterest.com",
		"pogo.com",
		"popurls.com",
		"postsecret.com",
		"realclearpolitics.com",
		"realtor.com",
		"reddit.com",
		"redfin.com",
		"sciencemag.org",
		"sho.com",
		"slashdot.com",
		"slashdot.org",
		"slickdeals.net",
		"stripgenerator.com",
		"techcrunch.com",
		"technorati.com",
		"techradar.com",
		"techrepublic.com",
		"textsfromlastnight.com",
		"thenextweb.com",
		"thepointsguy.com",
		"theverge.com",
		"tomshardware.com",
		"trulia.com",
		"tumblr.com",
		"twitter.com",
		"twitterfall.com",
		"typepad.com",
		"wearables.com",
		"wired.com",
		"woot.com",
		"xkcd.com",
		"yelp.com",
		"youtube.com",
		"zdnet.com",
		"zillow.com",
	}

	// Open or create /etc/hf_predef_distractors and store it in *os.File type
	fmt.Println(localFilePaths["pDistractors"])
	predefDistractorsFileObj, err := os.OpenFile(localFilePaths["pDistractors"], os.O_CREATE|os.O_WRONLY, 0666)
	if err != nil {
		log.Fatalln(err)
	}
	defer func() {
		if err := predefDistractorsFileObj.Close(); err != nil {
			log.Printf("Error closing file: %s\n", err)
		}
	}()

	// Initialize a new writer, range loop over distractors slice, and write
	// each distractor domain to /etc/hf_predef_distractors, line by line.
	dataWriter := bufio.NewWriter(predefDistractorsFileObj)
	for _, v := range distractors {

		_, _ = dataWriter.WriteString(v + "\n")
	}

	if err := dataWriter.Flush(); err != nil {
		log.Fatalln(err)
	}
}

func main() {

	parser := argparse.NewParser("go-hyperfocus", "Block distracting websites and hyperfocus on your work")

	improveCmd := parser.NewCommand("improve", "Improve focus, block the distractors")
	loseCmd := parser.NewCommand("lose", "Lose focus, unblock the distractors")

	breakCmd := parser.NewCommand("break", "Take a break for a number of minutes")
	minutesForBreak := breakCmd.Int("", "minutes", &argparse.Options{Help: "Number of minutes to break for."})

	listCmd := parser.NewCommand("list", "List the distractors defined in the block file (/etc/hf_distractors)")
	predefinedCmd := parser.NewCommand("predefined", "Add predefined set of distractors to /etc/hf_predef_distractors")

	if isRoot() {

		getLocalFilePaths()

		err := parser.Parse(os.Args)
		if err != nil {
			fmt.Print(parser.Usage(err))
		}

		if improveCmd.Happened() {
			improveFocus()
		} else if loseCmd.Happened() {
			loseFocus()
		} else if breakCmd.Happened() {
			takeBreak(*minutesForBreak)
		} else if listCmd.Happened() {
			listDistractors()
		} else if predefinedCmd.Happened() {
			addPredefinedDistractors()
		} else {
			fmt.Println("Enter a subcommand; use --help or -h for details.")
		}
	} else {
		fmt.Println("Super user not detected. This program requires root privileges; please re-run with sudo or become root.")
	}
}