From eddf6d694dc991e16f02f2b6e126912ad6d4c767 Mon Sep 17 00:00:00 2001
From: "Ajay K. Dhyani" <66765596+dhyanio@users.noreply.github.com>
Date: Thu, 5 Dec 2024 09:39:33 -0800
Subject: [PATCH] Add support for YAML configuration files (#6)

* config file can be YAML

* config file can be YAML

* config file can be YAML

* chore: update tailscalemanager
---
 README.md                      | 16 +++++++++++++++-
 src/TailscaleManager.hs        | 11 +++++------
 src/TailscaleManager/Config.hs | 31 +++++++++++++++++++++++++++----
 tailscale-manager.cabal        |  2 ++
 4 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index b831413..5c68b53 100644
--- a/README.md
+++ b/README.md
@@ -33,8 +33,9 @@ normally support App Connectors.
 
 ## Example
 
-Here is a sample config file, in JSON format:
+Create a configuration file in either JSON or YAML format. Here’s an example:
 
+`config.json`
 ```json
 {
   "routes": [
@@ -50,6 +51,19 @@ Here is a sample config file, in JSON format:
   ]
 }
 ```
+or `config.yaml`
+```yaml
+routes:
+  - "172.16.0.0/22"
+  - "192.168.0.0/24"
+hostRoutes:
+  - "special-hostname1.example"
+  - "special-hostname2.example"
+awsManagedPrefixLists:
+  - "pl-02761f4a40454a3c9"
+extraArgs:
+  - "--webclient"
+```
 
 Run tailscale-manager:
 
diff --git a/src/TailscaleManager.hs b/src/TailscaleManager.hs
index 430c0ea..04d3376 100644
--- a/src/TailscaleManager.hs
+++ b/src/TailscaleManager.hs
@@ -41,7 +41,7 @@ helpText = [r|Tailscale routes manager
 Dynamically resolves a list of hostRoutes to IP addresses, then tells tailscale
 to advertise them as /32 routes along with any normal CIDR routes.
 
-Config file example:
+Config file example (JSON):
 
 {
   "routes": [
@@ -50,7 +50,7 @@ Config file example:
   ],
   "hostRoutes": [
     "special-hostname1.example",
-    "special-hostname2.example",
+    "special-hostname2.example"
   ],
   "awsManagedPrefixLists": [
     "pl-02761f4a40454a3c9"
@@ -61,7 +61,7 @@ Config file example:
 tsManagerOptions :: Parser TailscaleManagerOptions
 tsManagerOptions =
   TailscaleManagerOptions
-  <$> argument str (metavar "<configfile.json>")
+  <$> argument str (metavar "<configfile>")
   <*> switch (long "dryrun"
               <> help "Dryrun mode")
   <*> strOption (long "tailscale"
@@ -143,7 +143,7 @@ runOnce options prevRoutes = do
       logDelay
       return prevRoutes
 
--- |Emit a log message describing the difference between old and new route sets.
+-- | Emit a log message describing the difference between old and new route sets.
 logDiff :: Set IPRange -> Set IPRange -> IO ()
 logDiff prevRoutes newRoutes = do
   logger <- myLogger
@@ -170,8 +170,7 @@ shrinkRatio :: Foldable t
             -> Double    -- ^ Shrink ratio
 shrinkRatio old new = 1 - (1 / (fromIntegral (length old) / fromIntegral (length new)))
 
--- |Do all the hostname resolution and concat the results with static routes from
--- the config.
+-- | Generate routes based on config, resolving hostnames and AWS-managed prefix lists.
 generateRoutes :: TSConfig -> IO (Set IPRange)
 generateRoutes config = do
   hostRoutes <- resolveHostnamesToRoutes (tsHostRoutes config)
diff --git a/src/TailscaleManager/Config.hs b/src/TailscaleManager/Config.hs
index 7b35126..eb27763 100644
--- a/src/TailscaleManager/Config.hs
+++ b/src/TailscaleManager/Config.hs
@@ -3,15 +3,14 @@
 module TailscaleManager.Config where
 
 import Data.Aeson
+import Data.Aeson (eitherDecodeFileStrict)
 import Data.Aeson.IP ()
 import Data.ByteString.Lazy qualified as LB
 import Data.IP (IPRange)
 import Data.Maybe (fromMaybe)
 import Data.Text
-
--- |Parse our config file.  May throw AesonException on failure.
-loadConfig :: FilePath -> IO TSConfig
-loadConfig fp = LB.readFile fp >>= throwDecode
+import System.FilePath (takeExtension)
+import Data.Yaml (decodeFileEither)
 
 -- |Config file schema
 data TSConfig
@@ -35,3 +34,27 @@ instance FromJSON TSConfig where
                      , tsExtraArgs = fromMaybe [] extraArgs
                      , tsAWSManagedPrefixLists = fromMaybe [] awsManagedPrefixLists
                      })
+
+-- | Load configuration from a file, detecting format based on file extension.
+loadConfig :: FilePath -> IO TSConfig
+loadConfig path = do
+  case takeExtension path of
+    ".json" -> loadConfigFromJSON path
+    ".yaml" -> loadConfigFromYAML path
+    _       -> error "Unsupported file format. Please use .json or .yaml."
+
+-- | Load configuration from a JSON file.
+loadConfigFromJSON :: FilePath -> IO TSConfig
+loadConfigFromJSON path = do
+  result <- eitherDecodeFileStrict path
+  case result of
+    Left err -> error $ "Failed to parse JSON config: " ++ err
+    Right config -> return config
+
+-- | Load configuration from a YAML file.
+loadConfigFromYAML :: FilePath -> IO TSConfig
+loadConfigFromYAML path = do
+  result <- decodeFileEither path
+  case result of
+    Left err -> error $ "Failed to parse YAML config: " ++ show err
+    Right config -> return config
diff --git a/tailscale-manager.cabal b/tailscale-manager.cabal
index 45bf47b..82f8a19 100644
--- a/tailscale-manager.cabal
+++ b/tailscale-manager.cabal
@@ -50,6 +50,8 @@ common deps
       , protolude ^>= 0.3.4
       , raw-strings-qq ^>= 1.1
       , text ^>= 2.0.2
+      , filepath ^>= 1.4.2
+      , yaml ^>= 0.11.8
 
 executable tailscale-manager
     import:           warnings, deps