{-# LANGUAGE LambdaCase 
           , DeriveGeneric
#-}
{-|
Module         : HGreet.Client
Description    : Simplified communication with the greetd daemon.
Copyright      : (c) Hazel (Vawlpe), 2022
License        : GPL-3.0-or-later
Maintainer     : vawlpe@gmail.com
Stability      : experimental
Portability    : Linux

To use this module, first get the path to the socket of the greetd daemon from the environment variable @GREETD_SOCK@. This requires the greetd daemon to be running.
Then you can communicate with greetd using the `withSocketDo` function of this module, passing it a callback function that will have direct acccess to the open socket.
For a simplified communication scheme you can implement a @handler@ and pass it to @handleResponse@ alongside the socket and command to run on successful authentication.
-}
module HGreet.Client (withSocketDo, send, recv, handleResponse, PromptResult(..)) where

import Data.Maybe (fromJust)
import Data.Functor ((<&>))
import Control.Exception (bracket_)
import Control.Concurrent (threadDelay)
import GHC.Generics (Generic)
import System.IO
import System.Exit
import qualified Data.ByteString.Lazy as BL
import qualified Control.Exception as E
import qualified Network.Socket as NS
import qualified Network.Socket.ByteString as NSB
import qualified HGreet.Packet as P ( Request(..), Response(..), AuthMessageType(..), ErrorType(..), encodeRequest, decodeResponse, decodeLen )

{-
  The `withSocketDo` function takes a to the socket of the greetd daemon and a callback function which will have access to the open socket.
  Within the callback function, you may use the send and recv functions and the HGreet.Packet module to communicate with the greetd daemon directly.
  Alternatively you can use the `handleResponse` function to implement the default login routine of greetd given a handler function.
  For examples on how to do both of these, see the "hagreety" package.
-}
withSocketDo :: String              -- ^ Path to the socket of the greetd daemon, usually found in the environment variable @GREETD_SOCK@
             -> (NS.Socket -> IO a) -- ^ Callback function that will have direct access to the open socket.
             -> IO a                -- ^ Result of the callback function as an IO action.
withSocketDo :: String -> (Socket -> IO a) -> IO a
withSocketDo String
path Socket -> IO a
client = do
    IO Socket -> (Socket -> IO ()) -> (Socket -> IO a) -> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
E.bracket (String -> IO Socket
open String
path) Socket -> IO ()
NS.close Socket -> IO a
client
  where
    open :: String -> IO Socket
open String
path = IO Socket
-> (Socket -> IO ()) -> (Socket -> IO Socket) -> IO Socket
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
E.bracketOnError (String -> IO Socket
forall p. p -> IO Socket
sockOpen String
path) Socket -> IO ()
NS.close ((Socket -> IO Socket) -> IO Socket)
-> (Socket -> IO Socket) -> IO Socket
forall a b. (a -> b) -> a -> b
$ \Socket
sock -> do
        Socket -> SockAddr -> IO ()
NS.connect Socket
sock (SockAddr -> IO ()) -> SockAddr -> IO ()
forall a b. (a -> b) -> a -> b
$ String -> SockAddr
NS.SockAddrUnix String
path
        Socket -> IO Socket
forall (m :: * -> *) a. Monad m => a -> m a
return Socket
sock
    sockOpen :: p -> IO Socket
sockOpen p
path = do
        Socket
sock <- Family -> SocketType -> ProtocolNumber -> IO Socket
NS.socket Family
NS.AF_UNIX SocketType
NS.Stream ProtocolNumber
NS.defaultProtocol
        Socket -> SocketOption -> Int -> IO ()
NS.setSocketOption Socket
sock SocketOption
NS.ReuseAddr Int
1
        Socket -> IO Socket
forall (m :: * -> *) a. Monad m => a -> m a
return Socket
sock

{-
 Send a `HGreet.Packet.Request` to the greetd daemon given an open socket.
 Usually used within the callback of a `withSocketDo` function to implement low level communication with the greetd daemon.
-}
send :: NS.Socket -- ^ Open socket to the greetd daemon, usually obtained from the callback of a `withSocketDo` function.
     -> P.Request -- ^ `HGreet.Packet.Request` to send to the greetd daemon.
     -> IO ()     -- ^ Empty IO action result.
send :: Socket -> Request -> IO ()
send Socket
sock Request
req = Socket -> ByteString -> IO ()
NSB.sendAll Socket
sock (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ Request -> ByteString
P.encodeRequest Request
req

{-
 Receive a `HGreet.Packet.Response` from the greetd daemon given an open socket.
 Usually used within the callback of a withSocketDo function to implement low level communication with the greetd daemon.
-}
recv :: NS.Socket     -- ^ Open socket to the greetd daemon, usually obtained from the callback of a `withSocketDo` function.
     -> IO P.Response -- ^ `HGreet.Packet.Response` received from the greetd daemon as an IO action.
recv :: Socket -> IO Response
recv Socket
sock = do
    ByteString
len <- Socket -> Int -> IO ByteString
NSB.recv Socket
sock Int
4
    let len' :: Int
len' = ByteString -> Int
P.decodeLen ( ByteString -> ByteString
BL.fromStrict ByteString
len ) :: Int
    ByteString
packet <- Socket -> Int -> IO ByteString
NSB.recv Socket
sock Int
len'
    Response -> IO Response
forall (m :: * -> *) a. Monad m => a -> m a
return (Response -> IO Response) -> Response -> IO Response
forall a b. (a -> b) -> a -> b
$ ByteString -> Response
P.decodeResponse ByteString
packet

{-
 Generic default login routine for greetd.
 Allows you to simply slap a handler function to deal with user input and have a working greeter with minimal effort.
 Currently kinda broken, will fix soon.
-}
handleResponse :: (P.Response -> IO PromptResult) -- ^ Handler function that will be called for every response from greetd.
               -> P.Response                      -- ^ Response from greetd that will be passed to the handler function.
               -> NS.Socket                       -- ^ Open socket to the greetd daemon, usually obtained from the callback of a `withSocketDo` function.
               -> [String]                        -- ^ List of strings to pass as the command to execute to start session after authentication.
               -> IO ()                           -- ^ Empty IO action result.
handleResponse :: (Response -> IO PromptResult)
-> Response -> Socket -> [String] -> IO ()
handleResponse Response -> IO PromptResult
handler Response
resp Socket
sock [String]
cmd =
    Response -> IO PromptResult
handler Response
resp IO PromptResult -> (PromptResult -> IO ()) -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        PromptResult
Error -> do
            Socket -> Request -> IO ()
send Socket
sock Request
P.CancelSession
            Int -> IO ()
threadDelay Int
1500000
            IO ()
forall a. IO a
exitFailure
        Auth String
msg -> do
            Socket -> Request -> IO ()
send Socket
sock (Request -> IO ()) -> Request -> IO ()
forall a b. (a -> b) -> a -> b
$ Maybe String -> Request
P.PostAuthMessageResponse (Maybe String -> Request) -> Maybe String -> Request
forall a b. (a -> b) -> a -> b
$ String -> Maybe String
forall a. a -> Maybe a
Just String
msg
            Response
rsp <- Socket -> IO Response
recv Socket
sock
            (Response -> IO PromptResult)
-> Response -> Socket -> [String] -> IO ()
handleResponse Response -> IO PromptResult
handler Response
rsp Socket
sock [String]
cmd
        PromptResult
Info -> do 
            Response
rsp <- Socket -> IO Response
recv Socket
sock
            (Response -> IO PromptResult)
-> Response -> Socket -> [String] -> IO ()
handleResponse Response -> IO PromptResult
handler Response
rsp Socket
sock [String]
cmd
        PromptResult
Success -> do
            Socket -> Request -> IO ()
send Socket
sock (Request -> IO ()) -> Request -> IO ()
forall a b. (a -> b) -> a -> b
$ [String] -> Request
P.StartSession [String]
cmd
            Int -> IO ()
threadDelay Int
500000
            IO ()
forall a. IO a
exitSuccess

{- 
  Prompt result type for handler functions that work with `handleResponse`.
-}
data PromptResult
    = Success         -- ^ Successful prompt result.
    | Error           -- ^ Error prompt result.
    | Info            -- ^ Info prompt result.
    | Auth String     -- ^ Auth prompt result.
    deriving ((forall x. PromptResult -> Rep PromptResult x)
-> (forall x. Rep PromptResult x -> PromptResult)
-> Generic PromptResult
forall x. Rep PromptResult x -> PromptResult
forall x. PromptResult -> Rep PromptResult x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep PromptResult x -> PromptResult
$cfrom :: forall x. PromptResult -> Rep PromptResult x
Generic, PromptResult -> PromptResult -> Bool
(PromptResult -> PromptResult -> Bool)
-> (PromptResult -> PromptResult -> Bool) -> Eq PromptResult
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: PromptResult -> PromptResult -> Bool
$c/= :: PromptResult -> PromptResult -> Bool
== :: PromptResult -> PromptResult -> Bool
$c== :: PromptResult -> PromptResult -> Bool
Eq, Int -> PromptResult -> ShowS
[PromptResult] -> ShowS
PromptResult -> String
(Int -> PromptResult -> ShowS)
-> (PromptResult -> String)
-> ([PromptResult] -> ShowS)
-> Show PromptResult
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PromptResult] -> ShowS
$cshowList :: [PromptResult] -> ShowS
show :: PromptResult -> String
$cshow :: PromptResult -> String
showsPrec :: Int -> PromptResult -> ShowS
$cshowsPrec :: Int -> PromptResult -> ShowS
Show)