{-# LANGUAGE CPP #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE OverloadedLists #-}

module Test.Sandwich.WebDriver.Internal.Capabilities.Extra (
  configureHeadlessCapabilities
  , configureDownloadCapabilities
  ) where

import Control.Monad.Catch (MonadMask)
import Control.Monad.IO.Class
import Control.Monad.IO.Unlift
import Control.Monad.Logger
import qualified Data.Aeson as A
import Data.Function
import qualified Data.List as L
import Data.Maybe
import Data.String.Interpolate
import qualified Data.Text as T
import qualified Data.Vector as V
import GHC.Stack
import Lens.Micro
import Lens.Micro.Aeson
import Test.Sandwich
import Test.Sandwich.WebDriver.Internal.Binaries.Chrome.Detect (detectChromeVersion)
import Test.Sandwich.WebDriver.Internal.Binaries.Chrome.Types (ChromeVersion(..))
import Test.Sandwich.WebDriver.Internal.Types
import qualified Test.WebDriver as W
import qualified Test.WebDriver.Firefox.Profile as FF


#if MIN_VERSION_aeson(2,0,0)
import qualified Data.Aeson.Key             as A
import qualified Data.Aeson.KeyMap          as HM
fromText :: T.Text -> A.Key
fromText :: Text -> Key
fromText = Text -> Key
A.fromText
#else
import qualified Data.HashMap.Strict        as HM
fromText :: T.Text -> T.Text
fromText = id
#endif


type Constraints m = (HasCallStack, MonadLogger m, MonadUnliftIO m, MonadMask m)

-- | Add headless configuration to the Chrome browser
configureHeadlessCapabilities :: (Constraints m) => WdOptions -> RunMode -> W.Capabilities -> m W.Capabilities
configureHeadlessCapabilities :: forall (m :: * -> *).
Constraints m =>
WdOptions -> RunMode -> Capabilities -> m Capabilities
configureHeadlessCapabilities WdOptions
_wdOptions (RunHeadless (HeadlessConfig {Maybe (Int, Int)
headlessResolution :: Maybe (Int, Int)
headlessResolution :: HeadlessConfig -> Maybe (Int, Int)
..})) caps :: Capabilities
caps@(W.Capabilities {browser :: Capabilities -> Browser
W.browser=browser :: Browser
browser@(W.Chrome {[String]
[ChromeExtension]
Maybe String
Object
chromeDriverVersion :: Maybe String
chromeBinary :: Maybe String
chromeOptions :: [String]
chromeExtensions :: [ChromeExtension]
chromeExperimentalOptions :: Object
chromeDriverVersion :: Browser -> Maybe String
chromeBinary :: Browser -> Maybe String
chromeOptions :: Browser -> [String]
chromeExtensions :: Browser -> [ChromeExtension]
chromeExperimentalOptions :: Browser -> Object
..})}) = do
  String
chromeBinaryPath <- case Maybe String
chromeBinary of
    Maybe String
Nothing -> String -> m String
forall (m :: * -> *) a. (HasCallStack, MonadIO m) => String -> m a
expectationFailure [i|Chrome capabilities didn't define chromeBinary in configureHeadlessCapabilities|]
    Just String
x -> String -> m String
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure String
x

  String
headlessArg <- IO (Either Text ChromeVersion) -> m (Either Text ChromeVersion)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (String -> IO (Either Text ChromeVersion)
detectChromeVersion String
chromeBinaryPath) m (Either Text ChromeVersion)
-> (Either Text ChromeVersion -> m String) -> m String
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Left Text
err -> do
      Text -> m ()
forall (m :: * -> *). (HasCallStack, MonadLogger m) => Text -> m ()
warn [i|Couldn't determine chrome version when configuring headless capabilities (err: #{err}); passing --headless|]
      String -> m String
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return String
"--headless"
    Right (ChromeVersion (Int
major, Int
_, Int
_, Int
_))
      -- See https://www.selenium.dev/blog/2023/headless-is-going-away/
      | Int
major Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
110 -> String -> m String
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return String
"--headless=new"
      | Bool
otherwise -> String -> m String
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return String
"--headless"

  let browser' :: Browser
browser' = Browser
browser { W.chromeOptions = headlessArg:resolution:chromeOptions }

  Capabilities -> m Capabilities
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Capabilities
caps { W.browser = browser' })

  where
    resolution :: String
resolution = [i|--window-size=#{w},#{h}|]
    (Int
w, Int
h) = (Int, Int) -> Maybe (Int, Int) -> (Int, Int)
forall a. a -> Maybe a -> a
fromMaybe (Int
1920, Int
1080) Maybe (Int, Int)
headlessResolution

-- | Add headless configuration to the Firefox capabilities
configureHeadlessCapabilities WdOptions
_ (RunHeadless (HeadlessConfig {})) caps :: Capabilities
caps@(W.Capabilities {browser :: Capabilities -> Browser
W.browser=(W.Firefox {}), additionalCaps :: Capabilities -> [Pair]
W.additionalCaps=[Pair]
ac}) = Capabilities -> m Capabilities
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Capabilities
caps { W.additionalCaps = additionalCaps })
  where
    additionalCaps :: [Pair]
additionalCaps = case (Pair -> Bool) -> [Pair] -> Maybe Int
forall a. (a -> Bool) -> [a] -> Maybe Int
L.findIndex (\Pair
x -> Pair -> Key
forall a b. (a, b) -> a
fst Pair
x Key -> Key -> Bool
forall a. Eq a => a -> a -> Bool
== Key
"moz:firefoxOptions") [Pair]
ac of
      Maybe Int
Nothing -> (Key
"moz:firefoxOptions", [Pair] -> Value
A.object [(Key
"args", Array -> Value
A.Array [Value
Item Array
"-headless"])]) Pair -> [Pair] -> [Pair]
forall a. a -> [a] -> [a]
: [Pair]
ac
      Just Int
i' -> let ffOptions' :: Value
ffOptions' = Pair -> Value
forall a b. (a, b) -> b
snd ([Pair]
ac [Pair] -> Int -> Pair
forall a. HasCallStack => [a] -> Int -> a
!! Int
i')
                                Value -> (Value -> Value) -> Value
forall a b. a -> (a -> b) -> b
& Text -> Value -> Value -> Value
ensureKeyExists Text
"args" (Array -> Value
A.Array [])
                                Value -> (Value -> Value) -> Value
forall a b. a -> (a -> b) -> b
& ((Key -> Traversal' Value Value
forall t. AsValue t => Key -> Traversal' t Value
key Key
"args" ((Value -> Identity Value) -> Value -> Identity Value)
-> ((Array -> Identity Array) -> Value -> Identity Value)
-> (Array -> Identity Array)
-> Value
-> Identity Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Array -> Identity Array) -> Value -> Identity Value
forall t. AsValue t => Traversal' t Array
Traversal' Value Array
_Array) ((Array -> Identity Array) -> Value -> Identity Value)
-> (Array -> Array) -> Value -> Value
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ Array -> Array
addHeadlessArg) in
        (Pair -> Pair -> Bool) -> [Pair] -> [Pair]
forall a. (a -> a -> Bool) -> [a] -> [a]
L.nubBy (\Pair
x Pair
y -> Pair -> Key
forall a b. (a, b) -> a
fst Pair
x Key -> Key -> Bool
forall a. Eq a => a -> a -> Bool
== Pair -> Key
forall a b. (a, b) -> a
fst Pair
y) ((Key
"moz:firefoxOptions", Value
ffOptions') Pair -> [Pair] -> [Pair]
forall a. a -> [a] -> [a]
: [Pair]
ac)

    ensureKeyExists :: T.Text -> A.Value -> A.Value -> A.Value
    ensureKeyExists :: Text -> Value -> Value -> Value
ensureKeyExists Text
key' Value
_ val :: Value
val@(A.Object (Key -> Object -> Maybe Value
forall v. Key -> KeyMap v -> Maybe v
HM.lookup (Text -> Key
fromText Text
key') -> Just Value
_)) = Value
val
    ensureKeyExists Text
key' Value
defaultVal (A.Object m :: Object
m@(Key -> Object -> Maybe Value
forall v. Key -> KeyMap v -> Maybe v
HM.lookup (Text -> Key
fromText Text
key') -> Maybe Value
Nothing)) = Object -> Value
A.Object (Key -> Value -> Object -> Object
forall v. Key -> v -> KeyMap v -> KeyMap v
HM.insert (Text -> Key
fromText Text
key') Value
defaultVal Object
m)
    ensureKeyExists Text
_ Value
_ Value
_ = String -> Value
forall a. HasCallStack => String -> a
error String
"Expected Object in ensureKeyExists"

    addHeadlessArg :: V.Vector A.Value -> V.Vector A.Value
    addHeadlessArg :: Array -> Array
addHeadlessArg Array
xs | (Text -> Value
A.String Text
"-headless") Value -> Array -> Bool
forall a. Eq a => a -> Vector a -> Bool
`V.elem` Array
xs = Array
xs
    addHeadlessArg Array
xs = (Text -> Value
A.String Text
"-headless") Value -> Array -> Array
forall a. a -> Vector a -> Vector a
`V.cons` Array
xs

configureHeadlessCapabilities WdOptions
_ (RunHeadless {}) Capabilities
browser = String -> m Capabilities
forall a. HasCallStack => String -> a
error [i|Headless mode not yet supported for browser '#{browser}'|]
configureHeadlessCapabilities WdOptions
_ RunMode
_ Capabilities
browser = Capabilities -> m Capabilities
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Capabilities
browser


-- | Configure download capabilities to set the download directory and disable prompts
-- (since you can't test download prompts using Selenium)
configureDownloadCapabilities :: (
  MonadIO m
  ) => [Char] -> W.Capabilities -> m W.Capabilities
configureDownloadCapabilities :: forall (m :: * -> *).
MonadIO m =>
String -> Capabilities -> m Capabilities
configureDownloadCapabilities String
downloadDir caps :: Capabilities
caps@(W.Capabilities {browser :: Capabilities -> Browser
W.browser=browser :: Browser
browser@(W.Firefox {Maybe Bool
Maybe String
Maybe (PreparedProfile Firefox)
LogLevel
ffProfile :: Maybe (PreparedProfile Firefox)
ffLogPref :: LogLevel
ffBinary :: Maybe String
ffAcceptInsecureCerts :: Maybe Bool
ffProfile :: Browser -> Maybe (PreparedProfile Firefox)
ffLogPref :: Browser -> LogLevel
ffBinary :: Browser -> Maybe String
ffAcceptInsecureCerts :: Browser -> Maybe Bool
..})}) = do
  PreparedProfile Firefox
profile <- case Maybe (PreparedProfile Firefox)
ffProfile of
    Just PreparedProfile Firefox
x -> PreparedProfile Firefox -> m (PreparedProfile Firefox)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure PreparedProfile Firefox
x
    Maybe (PreparedProfile Firefox)
Nothing -> IO (PreparedProfile Firefox) -> m (PreparedProfile Firefox)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (PreparedProfile Firefox) -> m (PreparedProfile Firefox))
-> IO (PreparedProfile Firefox) -> m (PreparedProfile Firefox)
forall a b. (a -> b) -> a -> b
$ Profile Firefox
FF.defaultProfile
      Profile Firefox
-> (Profile Firefox -> Profile Firefox) -> Profile Firefox
forall a b. a -> (a -> b) -> b
& Text -> Int -> Profile Firefox -> Profile Firefox
forall a b. ToPref a => Text -> a -> Profile b -> Profile b
FF.addPref Text
"browser.download.folderList" (Int
2 :: Int)
      Profile Firefox
-> (Profile Firefox -> Profile Firefox) -> Profile Firefox
forall a b. a -> (a -> b) -> b
& Text -> Bool -> Profile Firefox -> Profile Firefox
forall a b. ToPref a => Text -> a -> Profile b -> Profile b
FF.addPref Text
"browser.download.manager.showWhenStarting" Bool
False
      Profile Firefox
-> (Profile Firefox -> Profile Firefox) -> Profile Firefox
forall a b. a -> (a -> b) -> b
& Text -> String -> Profile Firefox -> Profile Firefox
forall a b. ToPref a => Text -> a -> Profile b -> Profile b
FF.addPref Text
"browser.download.dir" String
downloadDir
      Profile Firefox
-> (Profile Firefox -> Profile Firefox) -> Profile Firefox
forall a b. a -> (a -> b) -> b
& Text -> String -> Profile Firefox -> Profile Firefox
forall a b. ToPref a => Text -> a -> Profile b -> Profile b
FF.addPref Text
"browser.helperApps.neverAsk.saveToDisk" (String
"*" :: String)
      Profile Firefox
-> (Profile Firefox -> IO (PreparedProfile Firefox))
-> IO (PreparedProfile Firefox)
forall a b. a -> (a -> b) -> b
& Profile Firefox -> IO (PreparedProfile Firefox)
forall (m :: * -> *).
MonadBaseControl IO m =>
Profile Firefox -> m (PreparedProfile Firefox)
FF.prepareProfile

  Capabilities -> m Capabilities
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Capabilities
caps { W.browser = browser { W.ffProfile = Just profile } })
configureDownloadCapabilities String
downloadDir caps :: Capabilities
caps@(W.Capabilities {browser :: Capabilities -> Browser
W.browser=browser :: Browser
browser@(W.Chrome {[String]
[ChromeExtension]
Maybe String
Object
chromeDriverVersion :: Browser -> Maybe String
chromeBinary :: Browser -> Maybe String
chromeOptions :: Browser -> [String]
chromeExtensions :: Browser -> [ChromeExtension]
chromeExperimentalOptions :: Browser -> Object
chromeDriverVersion :: Maybe String
chromeBinary :: Maybe String
chromeOptions :: [String]
chromeExtensions :: [ChromeExtension]
chromeExperimentalOptions :: Object
..})}) = Capabilities -> m Capabilities
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Capabilities -> m Capabilities) -> Capabilities -> m Capabilities
forall a b. (a -> b) -> a -> b
$ Capabilities
caps { W.browser=browser' }
  where
    browser' :: Browser
browser' = Browser
browser { W.chromeExperimentalOptions = options }

    basePrefs :: A.Object
    basePrefs :: Object
basePrefs = case Key -> Object -> Maybe Value
forall v. Key -> KeyMap v -> Maybe v
HM.lookup Key
"prefs" Object
chromeExperimentalOptions of
      Just (A.Object Object
hm) -> Object
hm
      Just Value
x -> String -> Object
forall a. HasCallStack => String -> a
error [i|Expected chrome prefs to be object, got '#{x}'.|]
      Maybe Value
Nothing -> Object
forall a. Monoid a => a
mempty

    prefs :: A.Object
    prefs :: Object
prefs = Object
basePrefs
          Object -> (Object -> Object) -> Object
forall a b. a -> (a -> b) -> b
& ((Object -> Object) -> (Object -> Object) -> Object -> Object)
-> (Object -> Object) -> [Object -> Object] -> Object -> Object
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (Object -> Object) -> (Object -> Object) -> Object -> Object
forall b c a. (b -> c) -> (a -> b) -> a -> c
(.) Object -> Object
forall a. a -> a
id [Key -> Value -> Object -> Object
forall v. Key -> v -> KeyMap v -> KeyMap v
HM.insert Key
k Value
v | (Key
k, Value
v) <- [Pair]
downloadPrefs]

    options :: Object
options = Key -> Value -> Object -> Object
forall v. Key -> v -> KeyMap v -> KeyMap v
HM.insert Key
"prefs" (Object -> Value
A.Object Object
prefs) Object
chromeExperimentalOptions

    downloadPrefs :: [Pair]
downloadPrefs = [(Key
"profile.default_content_setting_values.automatic_downloads", Scientific -> Value
A.Number Scientific
1)
                    , (Key
"profile.content_settings.exceptions.automatic_downloads.*.setting", Scientific -> Value
A.Number Scientific
1)
                    , (Key
"download.prompt_for_download", Bool -> Value
A.Bool Bool
False)
                    , (Key
"download.directory_upgrade", Bool -> Value
A.Bool Bool
True)
                    , (Key
"download.default_directory", Text -> Value
A.String (String -> Text
T.pack String
downloadDir))]
configureDownloadCapabilities String
_ Capabilities
browser = Capabilities -> m Capabilities
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Capabilities
browser