{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeOperators #-}

module Test.Sandwich.Contexts.Kubernetes.Util.Images (
  dockerPullIfNecessary
  , isDockerImagePresent

  , podmanPullIfNecessary
  , isPodmanImagePresent

  , readImageName
  , readUncompressedImageName
  , imageLoadSpecToImageName
  ) where

import Control.Monad.IO.Unlift
import Control.Monad.Logger
import Data.Aeson as A
import qualified Data.ByteString.Lazy as BL
import Data.String.Interpolate
import qualified Data.Text as T
import qualified Data.Vector as V
import Relude
import Safe
import System.Exit
import System.FilePath
import Test.Sandwich
import Test.Sandwich.Contexts.Kubernetes.Types
import Test.Sandwich.Contexts.Kubernetes.Util.Aeson
import UnliftIO.Directory
import UnliftIO.Process
import UnliftIO.Temporary


-- * Docker

-- | Pull an image using Docker if it isn't already present.
-- Returns 'True' if a pull was done.
dockerPullIfNecessary :: (MonadUnliftIO m, MonadLoggerIO m) => Text -> ImagePullPolicy -> m Bool
dockerPullIfNecessary :: forall (m :: * -> *).
(MonadUnliftIO m, MonadLoggerIO m) =>
Text -> ImagePullPolicy -> m Bool
dockerPullIfNecessary = String -> Text -> ImagePullPolicy -> m Bool
forall (m :: * -> *).
(MonadUnliftIO m, MonadLoggerIO m) =>
String -> Text -> ImagePullPolicy -> m Bool
commonPullIfNecessary String
"docker"

isDockerImagePresent :: (MonadUnliftIO m, MonadLoggerIO m) => Text -> m Bool
isDockerImagePresent :: forall (m :: * -> *).
(MonadUnliftIO m, MonadLoggerIO m) =>
Text -> m Bool
isDockerImagePresent = String -> Text -> m Bool
forall (m :: * -> *).
(MonadUnliftIO m, MonadLoggerIO m) =>
String -> Text -> m Bool
isImagePresentCommon String
"docker"

-- * Podman

-- | Pull an image using Docker if it isn't already present.
-- Returns 'True' if a pull was done.
podmanPullIfNecessary :: (MonadUnliftIO m, MonadLoggerIO m) => Text -> ImagePullPolicy -> m Bool
podmanPullIfNecessary :: forall (m :: * -> *).
(MonadUnliftIO m, MonadLoggerIO m) =>
Text -> ImagePullPolicy -> m Bool
podmanPullIfNecessary = String -> Text -> ImagePullPolicy -> m Bool
forall (m :: * -> *).
(MonadUnliftIO m, MonadLoggerIO m) =>
String -> Text -> ImagePullPolicy -> m Bool
commonPullIfNecessary String
"podman"

isPodmanImagePresent :: (MonadUnliftIO m, MonadLoggerIO m) => Text -> m Bool
isPodmanImagePresent :: forall (m :: * -> *).
(MonadUnliftIO m, MonadLoggerIO m) =>
Text -> m Bool
isPodmanImagePresent = String -> Text -> m Bool
forall (m :: * -> *).
(MonadUnliftIO m, MonadLoggerIO m) =>
String -> Text -> m Bool
isImagePresentCommon String
"podman"

-- * Common

commonPullIfNecessary :: (MonadUnliftIO m, MonadLoggerIO m) => String -> Text -> ImagePullPolicy -> m Bool
commonPullIfNecessary :: forall (m :: * -> *).
(MonadUnliftIO m, MonadLoggerIO m) =>
String -> Text -> ImagePullPolicy -> m Bool
commonPullIfNecessary String
binary Text
image ImagePullPolicy
pullPolicy = String -> Text -> m Bool
forall (m :: * -> *).
(MonadUnliftIO m, MonadLoggerIO m) =>
String -> Text -> m Bool
isImagePresentCommon String
binary Text
image m Bool -> (Bool -> m Bool) -> m Bool
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  Bool
True ->
    if | ImagePullPolicy
pullPolicy ImagePullPolicy -> ImagePullPolicy -> Bool
forall a. Eq a => a -> a -> Bool
== ImagePullPolicy
Always -> m Bool
doPull
       | Bool
otherwise -> Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
  Bool
False ->
    if | ImagePullPolicy
pullPolicy ImagePullPolicy -> ImagePullPolicy -> Bool
forall a. Eq a => a -> a -> Bool
== ImagePullPolicy
Never -> String -> m Bool
forall (m :: * -> *) a. (HasCallStack, MonadIO m) => String -> m a
expectationFailure [i|Docker pull policy was "Never" but image wasn't present: '#{image}'|]
       | Bool
otherwise -> m Bool
doPull
  where
    doPull :: m Bool
doPull = do
      CreateProcess -> m ProcessHandle
forall (m :: * -> *).
(HasCallStack, MonadUnliftIO m, MonadLogger m) =>
CreateProcess -> m ProcessHandle
createProcessWithLogging (String -> [String] -> CreateProcess
proc String
binary [String
"pull", Text -> String
forall a. ToString a => a -> String
toString Text
image])
        m ProcessHandle -> (ProcessHandle -> m ExitCode) -> m ExitCode
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ProcessHandle -> m ExitCode
forall (m :: * -> *). MonadIO m => ProcessHandle -> m ExitCode
waitForProcess m ExitCode -> (ExitCode -> m ()) -> m ()
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (ExitCode -> ExitCode -> m ()
forall (m :: * -> *) a.
(HasCallStack, MonadIO m, Eq a, Show a) =>
a -> a -> m ()
`shouldBe` ExitCode
ExitSuccess)
      Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True

isImagePresentCommon :: (MonadUnliftIO m, MonadLoggerIO m) => String -> Text -> m Bool
isImagePresentCommon :: forall (m :: * -> *).
(MonadUnliftIO m, MonadLoggerIO m) =>
String -> Text -> m Bool
isImagePresentCommon String
binary Text
image = do
  CreateProcess -> m ProcessHandle
forall (m :: * -> *).
(HasCallStack, MonadUnliftIO m, MonadLogger m) =>
CreateProcess -> m ProcessHandle
createProcessWithLogging (String -> [String] -> CreateProcess
proc String
binary [String
"inspect", String
"--type=image", Text -> String
forall a. ToString a => a -> String
toString Text
image]) m ProcessHandle -> (ProcessHandle -> m ExitCode) -> m ExitCode
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ProcessHandle -> m ExitCode
forall (m :: * -> *). MonadIO m => ProcessHandle -> m ExitCode
waitForProcess m ExitCode -> (ExitCode -> m Bool) -> m Bool
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    ExitCode
ExitSuccess -> Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
    ExitFailure Int
_ -> Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False

-- * Image name reading

readImageName :: (HasCallStack, MonadUnliftIO m, MonadLogger m) => FilePath -> m Text
readImageName :: forall (m :: * -> *).
(HasCallStack, MonadUnliftIO m, MonadLogger m) =>
String -> m Text
readImageName String
path = String -> m Bool
forall (m :: * -> *). MonadIO m => String -> m Bool
doesDirectoryExist String
path m Bool -> (Bool -> m Text) -> m Text
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  Bool
True -> String -> m Text
forall (m :: * -> *). (HasCallStack, MonadIO m) => String -> m Text
readUncompressedImageName String
path
  Bool
False -> case String -> String
takeExtension String
path of
    String
".tar" -> m Text
extractFromTarball
    String
".gz" -> m Text
extractFromTarball
    String
_ -> String -> m Text
forall (m :: * -> *) a. (HasCallStack, MonadIO m) => String -> m a
expectationFailure [i|readImageName: unexpected extension in #{path}. Wanted .tar, .tar.gz, or uncompressed directory.|]
  where
    extractFromTarball :: m Text
extractFromTarball = do
      files <- CreateProcess -> String -> m String
forall (m :: * -> *).
(HasCallStack, MonadUnliftIO m, MonadLogger m) =>
CreateProcess -> String -> m String
readCreateProcessWithLogging (String -> [String] -> CreateProcess
proc String
"tar" [String
"tf", String
path]) String
""
      manifestFileName <- case headMay [t | t <- T.words (toText files), "manifest.json" `T.isInfixOf` t] of
        Just Text
f -> String -> m String
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> m String) -> String -> m String
forall a b. (a -> b) -> a -> b
$ Text -> String
forall a. ToString a => a -> String
toString (Text -> String) -> Text -> String
forall a b. (a -> b) -> a -> b
$ Text -> Text
T.strip Text
f
        Maybe Text
Nothing -> String -> m String
forall (m :: * -> *) a. (HasCallStack, MonadIO m) => String -> m a
expectationFailure [i|readImageName: couldn't find manifest file in #{path}|]

      withSystemTempDirectory "manifest.json" $ \String
dir -> do
        _ <- CreateProcess -> String -> m String
forall (m :: * -> *).
(HasCallStack, MonadUnliftIO m, MonadLogger m) =>
CreateProcess -> String -> m String
readCreateProcessWithLogging ((String -> [String] -> CreateProcess
proc String
"tar" [String
"xvf", String
path, String
manifestFileName]) { cwd = Just dir }) String
""
        liftIO (BL.readFile (dir </> "manifest.json")) >>= getImageNameFromManifestJson path

readUncompressedImageName :: (HasCallStack, MonadIO m) => FilePath -> m Text
readUncompressedImageName :: forall (m :: * -> *). (HasCallStack, MonadIO m) => String -> m Text
readUncompressedImageName String
path = IO ByteString -> m ByteString
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (String -> IO ByteString
BL.readFile (String
path String -> String -> String
</> String
"manifest.json")) m ByteString -> (ByteString -> m Text) -> m Text
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= String -> ByteString -> m Text
forall (m :: * -> *).
(HasCallStack, MonadIO m) =>
String -> ByteString -> m Text
getImageNameFromManifestJson String
path

getImageNameFromManifestJson :: (HasCallStack, MonadIO m) => FilePath -> LByteString -> m Text
getImageNameFromManifestJson :: forall (m :: * -> *).
(HasCallStack, MonadIO m) =>
String -> ByteString -> m Text
getImageNameFromManifestJson String
path ByteString
contents = do
  case ByteString -> Either String Value
forall a. FromJSON a => ByteString -> Either String a
A.eitherDecode ByteString
contents of
    Left String
err -> String -> m Text
forall (m :: * -> *) a. (HasCallStack, MonadIO m) => String -> m a
expectationFailure [i|Couldn't decode manifest.json: #{err}|]
    Right (A.Array Array
entries) -> case (Value -> [Text]) -> Array -> [Text]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Value -> [Text]
getRepoTags Array
entries of
      (Text
x:[Text]
_) -> Text -> m Text
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
x
      [] -> String -> m Text
forall (m :: * -> *) a. (HasCallStack, MonadIO m) => String -> m a
expectationFailure [i|Didn't find a repo tag for image at #{path}|]
    Right Value
x -> String -> m Text
forall (m :: * -> *) a. (HasCallStack, MonadIO m) => String -> m a
expectationFailure [i|Unexpected manifest.json format: #{x}|]

  where
    getRepoTags :: A.Value -> [Text]
    getRepoTags :: Value -> [Text]
getRepoTags (A.Object (Text -> Object -> Maybe Value
forall v. Text -> KeyMap v -> Maybe v
aesonLookup Text
"RepoTags" -> Just (A.Array Array
repoItems))) = [Text
t | A.String Text
t <- Array -> [Value]
forall a. Vector a -> [a]
V.toList Array
repoItems]
    getRepoTags Value
_ = []

imageLoadSpecToImageName :: (MonadUnliftIO m, MonadLogger m) => ImageLoadSpec -> m Text
imageLoadSpecToImageName :: forall (m :: * -> *).
(MonadUnliftIO m, MonadLogger m) =>
ImageLoadSpec -> m Text
imageLoadSpecToImageName (ImageLoadSpecTarball String
image) = String -> m Text
forall (m :: * -> *).
(HasCallStack, MonadUnliftIO m, MonadLogger m) =>
String -> m Text
readImageName String
image
imageLoadSpecToImageName (ImageLoadSpecDocker Text
image ImagePullPolicy
_) = Text -> m Text
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
image
imageLoadSpecToImageName (ImageLoadSpecPodman Text
image ImagePullPolicy
_) = Text -> m Text
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
image