1 files changed, 210 insertions, 0 deletions
diff --git a/tools/url-blocker/Main.hs b/tools/url-blocker/Main.hs
new file mode 100644
index 000000000000..34197981598c
--- /dev/null
+++ b/tools/url-blocker/Main.hs
@@ -0,0 +1,210 @@
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE DeriveGeneric #-}
+module Main
+  ( main
+  , getRules
+  , URL(..)
+  , Rule(..)
+  ) where
+
+--------------------------------------------------------------------------------
+-- Dependencies
+--------------------------------------------------------------------------------
+
+import qualified Data.Maybe as Maybe
+import qualified Data.Time.Clock as Clock
+import qualified Data.Time.Calendar as Calendar
+import qualified Data.Time.LocalTime as LocalTime
+import qualified Data.ByteString.Lazy as LazyByteString
+import qualified Data.Aeson as Aeson
+import qualified Data.Either.Combinators as Either
+import qualified Data.HashMap.Strict as HashMap
+import qualified Data.Text as Text
+import qualified Data.Text.IO as TextIO
+import qualified Data.Text.Read as TextRead
+import qualified Data.List as List
+
+import GHC.Generics
+import Data.Aeson ((.:))
+import Data.Text (Text)
+
+--------------------------------------------------------------------------------
+-- Types
+--------------------------------------------------------------------------------
+
+newtype URL = URL { getURL :: Text } deriving (Show, Eq, Generic)
+
+newtype IPAddress = IPAddress { getIPAddress :: Text } deriving (Show)
+
+newtype Domain = Domain { getDomain :: Text } deriving (Show)
+
+newtype Hour = Hour { getHour :: Int } deriving (Show, Eq, Generic)
+
+newtype Minute = Minute { getMinute :: Int } deriving (Show, Eq, Generic)
+
+data EtcHostsEntry = EtcHostsEntry { ip :: IPAddress
+                                   , domains :: [Domain]
+                                   } deriving (Show)
+
+-- | Write these in terms of your system's local time (i.e. `date`).
+data TimeSlot = TimeSlot { beg :: (Hour, Minute)
+                         , end :: (Hour, Minute)
+                         } deriving (Show, Eq, Generic)
+
+data Allowance = Allowance { day :: Calendar.DayOfWeek
+                           , timeslots :: [TimeSlot]
+                           } deriving (Show, Eq, Generic)
+
+data Rule = Rule { urls :: [URL]
+                 , allowed :: [Allowance]
+                 } deriving (Show, Eq, Generic)
+
+--------------------------------------------------------------------------------
+-- Instances
+--------------------------------------------------------------------------------
+
+instance Aeson.FromJSON TimeSlot where
+  parseJSON = Aeson.withText "timeslot" $ \x -> do
+    let [a, b] = Text.splitOn "-" x
+        [ah, am] = Text.splitOn ":" a
+        [bh, bm] = Text.splitOn ":" b
+    case extractTimeSlot ah am bh bm of
+      Left s  -> fail s
+      Right x -> pure x
+    where
+      extractTimeSlot :: Text -> Text -> Text -> Text -> Either String TimeSlot
+      extractTimeSlot ah am bh bm = do
+        (begh, _) <- TextRead.decimal ah
+        (begm, _) <- TextRead.decimal am
+        (endh, _) <- TextRead.decimal bh
+        (endm, _) <- TextRead.decimal bm
+        pure $ TimeSlot{ beg = (Hour begh, Minute begm)
+                       , end = (Hour endh, Minute endm)
+                       }
+
+instance Aeson.FromJSON Allowance where
+  parseJSON = Aeson.withObject "allowance" $ \x -> do
+    day <- x .: "day"
+    timeslots <- x .: "timeslots"
+    pure $ Allowance{day, timeslots}
+
+instance Aeson.FromJSON URL where
+  parseJSON = Aeson.withText "URL" $ \x -> do
+    pure $ URL { getURL = x }
+
+instance Aeson.FromJSON Rule where
+  parseJSON = Aeson.withObject "rule" $ \x -> do
+    urls <- x .: "urls"
+    allowed <- x .: "allowed"
+    pure Rule{urls, allowed}
+
+--------------------------------------------------------------------------------
+-- Functions
+--------------------------------------------------------------------------------
+
+-- | Pipe operator
+(|>) :: a -> (a -> b) -> b
+(|>) a f = f a
+infixl 1 |>
+
+-- | Returns True if the current time falls within any of the `timeslots`.
+isWithinTimeSlot :: LocalTime.LocalTime -> [TimeSlot] -> Bool
+isWithinTimeSlot date timeslots =
+  List.any withinTimeSlot timeslots
+  where
+    withinTimeSlot :: TimeSlot -> Bool
+    withinTimeSlot TimeSlot{ beg = (Hour ah, Minute am)
+                           , end = (Hour bh, Minute bm)
+                           } =
+      let LocalTime.TimeOfDay{LocalTime.todHour, LocalTime.todMin} =
+            LocalTime.localTimeOfDay date
+      in (todHour > ah) && (todMin > am) && (todHour < bh) && (todMin < bm)
+
+-- | Returns True if `day` is the same day as today.
+isToday :: LocalTime.LocalTime -> Calendar.DayOfWeek -> Bool
+isToday date day = today == day
+  where
+    today = Calendar.dayOfWeek (LocalTime.localDay date)
+
+-- | Returns True if a list of none of the `allowances` are valid.
+shouldBeBlocked :: LocalTime.LocalTime -> [Allowance] -> Bool
+shouldBeBlocked _ [] = True
+shouldBeBlocked date allowances = do
+  case filter (isToday date . day) allowances of
+    [Allowance{timeslots}] -> not $ isWithinTimeSlot date timeslots
+    [] -> True
+    -- Error when more than one rule per day
+    _  -> True
+
+-- | Maps an EtcHostsEntry to the line of text url-blocker will append to /etc/hosts.
+serializeEtcHostEntry :: EtcHostsEntry -> Text
+serializeEtcHostEntry EtcHostsEntry{ip, domains} =
+  (getIPAddress ip) <> "\t" <> (Text.unwords $ fmap getDomain domains)
+
+-- | Create an EtcHostsEntry mapping the URLs in `rule` to 127.0.0.1 if the
+-- URLs should be blocked.
+maybeBlockURL :: LocalTime.LocalTime -> Rule -> Maybe EtcHostsEntry
+maybeBlockURL date Rule{urls, allowed} =
+  if shouldBeBlocked date allowed then
+    Just $ EtcHostsEntry { ip = IPAddress "127.0.0.1"
+                        , domains = fmap (Domain . getURL) urls
+                        }
+  else
+    Nothing
+
+-- | Read and parse the rules.json file.
+-- TODO(wpcarro): Properly handle errors for file not found.
+-- TODO(wpcarro): Properly handle errors for parse failures.
+-- TODO(wpcarro): How can we resolve the $HOME directory when this is run as
+-- root?
+getRules :: IO [Rule]
+getRules = do
+  contents <- LazyByteString.readFile "/home/wpcarro/.config/url-blocker/rules.json"
+  let payload = Aeson.eitherDecode contents
+  pure $ Either.fromRight [] payload
+
+-- | Informational header added to /etc/hosts before the entries that
+-- url-blocker adds.
+urlBlockerHeader :: Text
+urlBlockerHeader =
+  Text.unlines [ "################################################################################"
+               , "# Added by url-blocker."
+               , "#"
+               , "# Warning: url-blocker will remove anything that you add beneath this header."
+               , "################################################################################"
+               ]
+
+-- | Removes all entries that url-blocker may have added to /etc/hosts.
+removeURLBlockerEntries :: Text -> Text
+removeURLBlockerEntries etcHosts =
+  case Text.breakOn urlBlockerHeader etcHosts of
+    (etcHosts', _) -> etcHosts'
+
+-- | Appends the newly created `entries` to `etcHosts`.
+addURLBlockerEntries :: Text -> Text -> Text
+addURLBlockerEntries entries etcHosts =
+  Text.unlines [ etcHosts
+               , urlBlockerHeader
+               , entries
+               ]
+
+-- | This script reads the current /etc/hosts, removes any entries that
+-- url-blocker may have added in a previous run, and adds new entries to block
+-- URLs according to the rules.json file.
+main :: IO ()
+main = do
+  rules <- getRules
+  tz <- LocalTime.getCurrentTimeZone
+  ct <- Clock.getCurrentTime
+  let date = LocalTime.utcToLocalTime tz ct
+      entries = rules
+                |> fmap (maybeBlockURL date)
+                |> Maybe.catMaybes
+                |> fmap serializeEtcHostEntry
+                |> Text.unlines
+  existingEtcHosts <- TextIO.readFile "/etc/hosts"
+  existingEtcHosts
+    |> removeURLBlockerEntries
+    |> addURLBlockerEntries entries
+    |> \x -> writeFile "/etc/hosts" (Text.unpack x)