let seconds = 5
  hPutStrLn stderr $ "Waiting for " ++ show seconds ++  " seconds for " ++ domain
  threadDelay (seconds * 1000000)

      let seconds = 5
      hPutStrLn stderr $ "Waiting for " ++ show seconds ++  " seconds for " ++ domain
      threadDelay (seconds * 1000000)

{-# LANGUAGE OverloadedStrings #-}
module Env
  ( Env(..)
  , loadEnv
  ) where
import System.IO
import Control.Exception (catch, IOException)
import Data.List (isPrefixOf)
import Data.Maybe (fromMaybe)
import Data.Char (toUpper)
import Text.Read (readMaybe)
import Udp2Raw (RawMode(..))  -- Import the RawMode type (ICMP, UDP, FakeTCP)
-------------------------------------------------------------------------------
-- | Environment configuration
-------------------------------------------------------------------------------
data Env = Env
  { ip      :: String
  , rawMode :: RawMode
  } deriving (Show, Eq)
-------------------------------------------------------------------------------
-- | Load and parse the .env file
-------------------------------------------------------------------------------
loadEnv :: FilePath -> IO (Either String Env)
loadEnv path = catch (do
    content <- readFile path
    let kvs = parseEnv content
        mIP = lookup "SERVER_IP" kvs
        mMode = lookup "RAW_MODE" kvs
        mPort = lookup "PORT" kvs
    case (mIP, mMode) of
      (Just ipStr, Just modeStr) ->
        case parseRawMode modeStr mPort of
          Just raw -> pure $ Right $ Env ipStr raw
          Nothing  -> pure $ Left $ "Invalid RAW_MODE or missing PORT for mode: " ++ modeStr
      _ -> pure $ Left "Missing SERVER_IP or RAW_MODE in .env"
  ) handleReadError
  where
    handleReadError :: IOException -> IO (Either String Env)
    handleReadError e = pure $ Left $ "Error reading .env: " ++ show e
-------------------------------------------------------------------------------
-- | Parse RAW_MODE + optional PORT into RawMode
-------------------------------------------------------------------------------
parseRawMode :: String -> Maybe String -> Maybe RawMode
parseRawMode modeStr mPort =
  let mode = map toUpper modeStr
      port = mPort >>= readMaybe
  in case mode of
       "ICMP"    -> Just ICMP
       "UDP"     -> UDP <$> port         -- mandatory port
       "FAKETCP" -> FakeTCP <$> port     -- mandatory port
       "TCP"     -> FakeTCP <$> port     -- alias for FakeTCP
       _         -> Nothing
-------------------------------------------------------------------------------
-- | Parse .env content into key-value pairs
-------------------------------------------------------------------------------
parseEnv :: String -> [(String, String)]
parseEnv = mapMaybe parseLine . lines
  where
    parseLine line =
      case break (== '=') (trim line) of
        (key, '=':value)
          | not (null key) && not (isComment key) ->
              Just (trim key, trim value)
        _ -> Nothing
    isComment ('#':_) = True
    isComment _       = False
    trim = f . f
      where f = reverse . dropWhile (`elem` [' ', '\t', '\r'])
-- | Helper: mapMaybe for base < 4.8
mapMaybe :: (a -> Maybe b) -> [a] -> [b]
mapMaybe f = foldr (\x acc -> case f x of
                               Just y  -> y:acc
                               Nothing -> acc) []

{-# LANGUAGE OverloadedStrings #-}
module ConfigParserUtils (getYamlConfigEntry) where
import ConfigParser (ConfigEntry, parseConfigFile)
import Data.List (find)
import Data.Maybe (listToMaybe)
import ConfigParser (ConfigEntry, parseConfigFile, name)  -- add `name`
import qualified Data.Text as T
-- | Parse YAML file and return either:
--   * The entry matching the given name, or
--   * The first entry if no name is provided
getYamlConfigEntry :: FilePath          -- ^ YAML config file path
                   -> Maybe String      -- ^ Optional name
                   -> IO (Either String ConfigEntry)
getYamlConfigEntry file mName = do
  parseResult <- parseConfigFile file
  case parseResult of
    Left err -> pure $ Left err
    Right cfg -> case mName of
      Just cfgNameStr -> case find (\e -> T.unpack (name e) == cfgNameStr) cfg of
                     Just entry -> pure $ Right entry
                     Nothing -> pure $ Left $ "No config entry named: " ++ cfgNameStr
      Nothing -> case listToMaybe cfg of
                   Just entry -> pure $ Right entry
                   Nothing -> pure $ Left "YAML config is empty"

#!/bin/sh
SOCK=/run/tproxy/socket
if [[ ! -S $SOCK ]]; then
  echo "tproxy service not running"
  exit 1
fi
if [[ $# -lt 1 ]]; then
  echo "usage: $0 {start|stop|status|exit} [name]"
  exit 1
fi
# Send all arguments joined by spaces
echo "$*" | nc -U "$SOCK"

      Env
      ConfigParserUtils

  - Env
  - ConfigParserUtils

import ConfigParser      -- Import the module we just created

import ConfigParser
import ConfigParserUtils (getYamlConfigEntry)
import Env

import Control.Concurrent (forkIO, MVar, newEmptyMVar, putMVar, takeMVar, killThread, forkFinally, isEmptyMVar)
import Control.Exception (catch, try, SomeException, IOException)

import Control.Concurrent
    ( forkIO, MVar, newEmptyMVar, putMVar, takeMVar, killThread, isEmptyMVar )
import Control.Exception (catch, IOException)

import System.Exit (exitFailure)
import System.IO (hPutStrLn, stderr, IOMode( ReadWriteMode ), hSetBuffering, BufferMode( LineBuffering ), hGetLine, hClose, stdout, BufferMode( NoBuffering ))
import Control.Monad (forM_, unless, forever, void, when)

import System.IO
    ( hPutStrLn, stderr, hSetBuffering, BufferMode(LineBuffering, NoBuffering)
    , hGetLine, hClose, stdout, IOMode(ReadWriteMode), Handle )

import Data.Maybe (fromMaybe)
import System.IO.Error (isDoesNotExistError)

import System.Directory (removeFile)
import Data.Maybe (fromMaybe, listToMaybe)
import Data.List (find)

import System.Exit (exitSuccess)
import System.Directory (removeFile)

  hSetBuffering stdout NoBuffering

    hSetBuffering stdout NoBuffering

  stopFlag <- newEmptyMVar -- Create a flag that the signal handler can set

    stopFlag <- newEmptyMVar
    installHandler sigTERM (Catch $ putMVar stopFlag ()) Nothing

  -- Install signal handlers for SIGTERM (systemd stop) and SIGINT (Ctrl+C)
  installHandler sigTERM (Catch $ putMVar stopFlag ()) Nothing
  -- 1. Get the configuration file path
  args <- getArgs
  let configFile = case args of
                     (f:_) -> f
                     [] -> "servers.yaml"

    args <- getArgs
    let configFile = case args of
                        (f:_) -> f
                        []    -> "servers.yaml"

  -- 2. Call the exported function from the ConfigParser module
  parseResult <- ConfigParser.parseConfigFile configFile

    -- Try to load .env first
    eEnv <- Env.loadEnv "/run/udp2raw/.env"
    case eEnv of
        Right envVal -> do
            putStrLn $ "Loaded environment: " ++ show envVal
            start (Env.ip envVal) (Env.rawMode envVal)
        Left _ -> do
            -- Fallback to YAML
            yamlCfg <- parseConfigFile configFile
            case yamlCfg of
                Right (entry:_) -> do
                    ip <- waitForIP (T.unpack $ nameserver entry) (T.unpack $ domain entry)
                    let raw = fromMaybe ICMP (textToRawMode (raw_mode entry) (port entry))
                    writeEnv ip (T.unpack $ udp2raw_password entry) (T.unpack $ udpspeeder_password entry) raw
                    start ip raw
                Right [] -> putStrLn "YAML configuration is empty"
                Left err -> hPutStrLn stderr $ "Failed to load .env and YAML: " ++ err

  case parseResult of
    -- Handle failure (The Left String contains the formatted error)
    Left errorMessage -> do
      hPutStrLn stderr errorMessage

    -- Clean up stale socket
    catch (removeFile socketPath) (\(_ :: IOException) -> pure ())
    sock <- socket AF_UNIX Stream 0
    bind sock (SockAddrUnix socketPath)
    setFileMode socketPath 0o775
    listen sock 10

    -- Handle success
    Right cfg -> do
      putStrLn $ "Successfully loaded configuration from " ++ configFile ++ ":"

    let envVal = either (const Nothing) Just eEnv

      -- Accessing the first entry's nameserver
      case cfg of
        (entry:_) -> do
          putStrLn $ "The server is " ++ show (name entry)
          ip <- waitForIP (T.unpack $ nameserver entry) (T.unpack $ domain entry)
          let raw = fromMaybe ICMP (textToRawMode (raw_mode entry) (port entry))
          let udp2raw_pwd = T.unpack $ udp2raw_password entry
          let udpspeeder_pwd = T.unpack $ udpspeeder_password entry
          writeEnv ip udp2raw_pwd udpspeeder_pwd raw
          start ip raw

    acceptThread <- forkIO $ acceptLoop sock stopFlag configFile envVal

          -- Clean up any existing socket file
          catch (removeFile socketPath)
                (\(e :: IOException) ->
                  if isDoesNotExistError e then pure () else ioError e)

    -- Wait for SIGTERM
    takeMVar stopFlag
    putStrLn "Received SIGTERM, shutting down gracefully..."
    killThread acceptThread
    close sock
    removeFile socketPath
    -- Stop proxy on shutdown
    case eEnv of
        Right envVal -> stop (Env.ip envVal) (Env.rawMode envVal)
        Left _ -> pure ()

          -- Create and bind the UNIX domain socket
          sock <- socket AF_UNIX Stream 0
          bind sock (SockAddrUnix socketPath)
          setFileMode socketPath 0o775
          listen sock 10

            -- Run the accept loop in a separate thread
          acceptThread <- forkIO $
            acceptLoop sock stopFlag (\conn -> handleClient conn ip raw)

-- | Accept loop
acceptLoop :: Socket -> MVar () -> FilePath -> Maybe Env.Env -> IO ()
acceptLoop sock stopFlag configFile envFallback = loop
  where
    loop = do
      empty <- isEmptyMVar stopFlag
      if empty
        then do
          (conn, _) <- accept sock
          forkIO $ handleClient conn configFile envFallback
          loop
        else pure ()

          -- Wait until SIGTERM
          takeMVar stopFlag
          putStrLn "Received SIGTERM, shutting down gracefully..."

-- | Handle a single client connection
-- | Handle a single client connection
handleClient :: Socket -> FilePath -> Maybe Env.Env -> IO ()
handleClient sock configFile envFallback = do
    h <- socketToHandle sock ReadWriteMode
    hSetBuffering h LineBuffering
    cmdLine <- hGetLine h
    let cmdWords = words cmdLine
    case cmdWords of
        ("start":nameParts) -> do
            let mName = if null nameParts then Nothing else Just (unwords nameParts)
            hPutStrLn h $ "Starting " ++ (unwords nameParts)

          --clear the iptables
          stop ip raw

            -- Parse YAML config
            yamlCfg <- parseConfigFile configFile
            mEntry <- case yamlCfg of
                Right cfg -> case mName of
                    -- Select entry by name
                    Just cfgNameStr ->
                      case find (\e -> T.strip (name e) == T.pack cfgNameStr) cfg of
                          Just e  -> do
                            hPutStrLn h $ "Found " ++ cfgNameStr
                            pure (Just e)
                          Nothing -> do
                              hPutStrLn h $ "No config entry named: " ++ cfgNameStr
                              pure Nothing
                    -- No name provided, pick the first entry
                    Nothing -> pure $ listToMaybe cfg
                Left _ -> pure Nothing
            (ipAddr, rawModeVal) <- case mEntry of
                Just entry -> do
                    ip <- waitForIP (T.unpack $ nameserver entry) (T.unpack $ domain entry)
                    let raw = fromMaybe ICMP (textToRawMode (raw_mode entry) (port entry))
                    writeEnv ip (T.unpack $ udp2raw_password entry) (T.unpack $ udpspeeder_password entry) raw
                    pure (ip, raw)
                Nothing -> case envFallback of
                    Just envVal -> pure (Env.ip envVal, Env.rawMode envVal)
                    Nothing -> do
                        hPutStrLn h "No matching config entry and no .env"
                        pure ("", ICMP)

          -- Stop accepting, clean up
          killThread acceptThread
          close sock
          removeFile socketPath
        [] -> putStrLn "Configuration file was empty."

            -- Start proxy with resolved IP and RawMode
            start ipAddr rawModeVal
            hPutStrLn h "started"
        ["stop"] -> do
            case envFallback of
                Just envVal -> do
                    stop (Env.ip envVal) (Env.rawMode envVal)
                    hPutStrLn h "stopped"
                Nothing -> hPutStrLn h "No .env loaded, cannot stop"
        _ -> hPutStrLn h "unknown command"
    hClose h

-- | Loop accepting connections until stopFlag is set
acceptLoop :: Socket -> MVar () -> (Socket -> IO ()) -> IO ()
acceptLoop sock stopFlag handleConn = do
  let loop = do
        stillRunning <- isEmptyMVar stopFlag
        when stillRunning $ do
          (conn, _) <- accept sock
          void $ forkFinally (handleConn conn) (\_ -> close conn)
          loop
  loop

handleClient :: Socket
                -> String   -- ^ server IP (resolved at runtime)
                -> RawMode
                -> IO ()
handleClient conn ip raw = do
  h <- socketToHandle conn ReadWriteMode
  hSetBuffering h LineBuffering
  cmd <- hGetLine h
  case cmd of
    "start" -> do
      putStrLn "Enabling global proxy"
      start ip raw
      hPutStrLn h "started"
    "stop" -> do
      putStrLn "Disabling global proxy"
      stop ip raw
      hPutStrLn h "stopped"
    _ -> hPutStrLn h "unknown command"
  hClose h

start ::  String       -- ^ server IP (resolved at runtime)
          -> RawMode
          -> IO ()

-- | Start the proxy
start :: String -> RawMode -> IO ()

  applyRules
  case raw of
    ICMP     -> applyICMP ip
    UDP p    -> applyUDP ip p
    FakeTCP p-> applyFakeTCP ip p

    applyRules
    case raw of
        ICMP      -> applyICMP ip
        UDP p     -> applyUDP ip p
        FakeTCP p -> applyFakeTCP ip p

stop ::   String       -- ^ server IP (resolved at runtime)
          -> RawMode
          -> IO ()

-- | Stop the proxy
stop :: String -> RawMode -> IO ()

  clearRules
  case raw of
      ICMP     -> do clearICMP ip
      UDP p    -> do clearUDP ip p
      FakeTCP p-> do clearFakeTCP ip p

    case raw of
        ICMP      -> clearICMP ip
        UDP p     -> clearUDP ip p
        FakeTCP p -> clearFakeTCP ip p