{-
 - This Source Code Form is subject to the terms of the Mozilla Public
 - License, v. 2.0. If a copy of the MPL was not distributed with this
 - file, You can obtain one at https://mozilla.org/MPL/2.0/.
 -}
{-# LANGUAGE OverloadedStrings #-}

module Test.System.Win32.Dpapi (dpapiTests) where

import Data.Base64.Types (extractBase64)
import Data.ByteString (toStrict)
import Data.ByteString.Base64 (decodeBase64Untyped, encodeBase64)
import Data.ByteString.Lazy (ByteString, dropWhileEnd)
import Data.Char (chr, isSpace)
import Data.Text (unpack)
import System.Process.Typed
  ( ExitCode (ExitSuccess),
    nullStream,
    proc,
    readProcessStdout,
    setStderr,
  )
import System.Win32.Dpapi
  ( DataProtectionScope (CurrentUser, LocalMachine),
    cryptProtectData,
    cryptUnprotectData,
    cryptUnprotectDataCheck,
  )
import Test.HUnit

dpapiTests :: Test
dpapiTests = "System.Win32.Dpapi" ~: [testProtectCurrentUser, testProtectLocalMachine, testUnprotectCurrentUser, testUnprotectLocalMachine, testUnprotectLocalMachineMaybe, testUnprotectError]

testProtectCurrentUser :: Test
testProtectCurrentUser =
  "testProtectCurrentUser" ~: do
    protected <- cryptProtectData "fizz" Nothing CurrentUser
    (exitCode, out) <- readProcessStdout $ setStderr nullStream $ proc "powershell" ["-c", "Add-Type -AssemblyName System.Security; [System.Text.Encoding]::UTF8.GetString([System.Security.Cryptography.ProtectedData]::Unprotect([System.Convert]::FromBase64String(\"" <> unpack (extractBase64 $ encodeBase64 protected) <> "\"), $null, [System.Security.Cryptography.DataProtectionScope]::CurrentUser))"]
    assertEqual "powershell exit code" ExitSuccess exitCode
    let out' = trimEndBs out
    out' @?= "fizz"

testProtectLocalMachine :: Test
testProtectLocalMachine =
  "testProtectLocalMachine" ~: do
    protected <- cryptProtectData "buzz" Nothing LocalMachine
    (exitCode, out) <- readProcessStdout $ setStderr nullStream $ proc "powershell" ["-c", "Add-Type -AssemblyName System.Security; [System.Text.Encoding]::UTF8.GetString([System.Security.Cryptography.ProtectedData]::Unprotect([System.Convert]::FromBase64String(\"" <> unpack (extractBase64 $ encodeBase64 protected) <> "\"), $null, [System.Security.Cryptography.DataProtectionScope]::LocalMachine))"]
    assertEqual "powershell exit code" ExitSuccess exitCode
    let out' = trimEndBs out
    out' @?= "buzz"

testUnprotectCurrentUser :: Test
testUnprotectCurrentUser =
  "testUnprotectCurrentUser" ~: do
    (exitCode, out) <- readProcessStdout $ setStderr nullStream $ proc "powershell" ["-c", "Add-Type -AssemblyName System.Security; [System.Convert]::ToBase64String([System.Security.Cryptography.ProtectedData]::Protect([System.Text.Encoding]::UTF8.GetBytes(\"foo\"), $null, [System.Security.Cryptography.DataProtectionScope]::CurrentUser))"]
    assertEqual "powershell exit code" ExitSuccess exitCode
    let out' = toStrict $ trimEndBs out -- Powershell outputs \r\n even when it's not wanted
    case decodeBase64Untyped out' of
      l@(Left _) ->
        Right out' @?= l
      Right bs -> do
        result <- cryptUnprotectData bs Nothing CurrentUser
        result @?= "foo"

testUnprotectLocalMachine :: Test
testUnprotectLocalMachine =
  "testUnprotectLocalMachine" ~: do
    (exitCode, out) <- readProcessStdout $ setStderr nullStream $ proc "powershell" ["-c", "Add-Type -AssemblyName System.Security; [System.Convert]::ToBase64String([System.Security.Cryptography.ProtectedData]::Protect([System.Text.Encoding]::UTF8.GetBytes(\"bar\"), $null, [System.Security.Cryptography.DataProtectionScope]::LocalMachine))"]
    assertEqual "powershell exit code" ExitSuccess exitCode
    let out' = toStrict $ trimEndBs out -- Powershell outputs \r\n even when it's not wanted
    case decodeBase64Untyped out' of
      l@(Left _) ->
        Right out' @?= l
      Right bs -> do
        result <- cryptUnprotectData bs Nothing LocalMachine
        result @?= "bar"

testUnprotectLocalMachineMaybe :: Test
testUnprotectLocalMachineMaybe =
  "testUnprotectLocalMachine" ~: do
    (exitCode, out) <- readProcessStdout $ setStderr nullStream $ proc "powershell" ["-c", "Add-Type -AssemblyName System.Security; [System.Convert]::ToBase64String([System.Security.Cryptography.ProtectedData]::Protect([System.Text.Encoding]::UTF8.GetBytes(\"bar\"), $null, [System.Security.Cryptography.DataProtectionScope]::LocalMachine))"]
    assertEqual "powershell exit code" ExitSuccess exitCode
    let out' = toStrict $ trimEndBs out -- Powershell outputs \r\n even when it's not wanted
    case decodeBase64Untyped out' of
      l@(Left _) ->
        Right out' @?= l
      Right bs -> do
        result <- cryptUnprotectDataCheck bs Nothing LocalMachine
        result @?= Just "bar"

testUnprotectError :: Test
testUnprotectError =
  "testUnprotectError" ~: do
    result <- cryptUnprotectDataCheck ciphertext Nothing CurrentUser
    result @?= Nothing
  where
    ciphertext = case decodeBase64Untyped "AQAAANCMnd8BFdERjHoAwE/Cl+sBAAAA8Zcdtib+QkybE9Tbcq0/EwQAAAACAAAAAAAQZgAAAAEAACAAAAB5qoj4OynizB5iEVZvAtyEZlVVn2RjP4SdHz5dOnIjYAAAAAAOgAAAAAIAACAAAAAGKRerRiwIK/xPPH8NtVhAxU9dg2qfcNsxLg+wtYml8RAAAADinEDleLmnP+qGkSpqgeZ9QAAAAAmaXs8XjSezA4qSWIFL3HD2bqd757HNPfXfuBz4eVYtfRawfRpFXFTmfIh+2nOzf1N2IYH44fBK7VrW+ITOZ74=" of
      (Right c) -> c
      (Left _) -> error "Programming error! Base64 text wasn't!"

trimEndBs :: ByteString -> ByteString
trimEndBs = dropWhileEnd (isSpace . chr . fromIntegral)