{-# LANGUAGE AllowAmbiguousTypes #-}

-- | This module provides an automated attack to try and redirect outputs to a
-- certain target with a similar datum type.
module Cooked.Attack.DatumHijacking
  ( redirectOutputTweakAny,
    datumHijackingAttackAny,
    datumHijackingAttack,
    redirectOutputTweakAll,
    datumHijackingAttackAll,
    DatumHijackingLbl (..),
  )
where

import Control.Monad
import Cooked.Pretty.Class
import Cooked.Skeleton
import Cooked.Tweak
import Data.Maybe
import Optics.Core
import Plutus.Script.Utils.Address qualified as Script
import PlutusLedgerApi.V3 qualified as Api
import Prettyprinter ((<+>))

-- | Redirects some outputs from one owner to another owner, which can be of
-- different types.
redirectOutputTweakAll ::
  forall owner owner' m.
  (MonadTweak m, OwnerConstrs owner, OwnerConstrs owner') =>
  -- | Return 'Just' the new owner, or 'Nothing' if you want to leave this
  -- output unchanged.
  (TxSkelOut -> Maybe owner') ->
  -- | The redirection described by the previous argument might apply to more
  -- than one of the outputs of the transaction. Use this predicate to select
  -- which of the redirectable outputs to actually redirect. We count the
  -- redirectable outputs from the left to the right, starting with zero.
  (Integer -> Bool) ->
  -- | Returns the list of outputs it redirected (as they were
  -- before the modification), in the order in which they occurred on the original
  -- transaction.
  m [TxSkelOut]
redirectOutputTweakAll :: forall owner owner' (m :: * -> *).
(MonadTweak m, OwnerConstrs owner, OwnerConstrs owner') =>
(TxSkelOut -> Maybe owner') -> (Integer -> Bool) -> m [TxSkelOut]
redirectOutputTweakAll TxSkelOut -> Maybe owner'
outputPred Integer -> Bool
indexPred = do
  [TxSkelOut]
outputs <- Optic' A_Lens NoIx TxSkel [TxSkelOut] -> m [TxSkelOut]
forall (m :: * -> *) k (is :: IxList) a.
(MonadTweak m, Is k A_Getter) =>
Optic' k is TxSkel a -> m a
viewTweak Optic' A_Lens NoIx TxSkel [TxSkelOut]
txSkelOutsL
  let ([Maybe TxSkelOut]
changed, [TxSkelOut]
newOutputs) = [(Maybe TxSkelOut, TxSkelOut)] -> ([Maybe TxSkelOut], [TxSkelOut])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Maybe TxSkelOut, TxSkelOut)]
 -> ([Maybe TxSkelOut], [TxSkelOut]))
-> [(Maybe TxSkelOut, TxSkelOut)]
-> ([Maybe TxSkelOut], [TxSkelOut])
forall a b. (a -> b) -> a -> b
$ [TxSkelOut] -> Integer -> [(Maybe TxSkelOut, TxSkelOut)]
go [TxSkelOut]
outputs Integer
0
  Optic' A_Lens NoIx TxSkel [TxSkelOut] -> [TxSkelOut] -> m ()
forall (m :: * -> *) k (is :: IxList) a.
(MonadTweak m, Is k A_Setter) =>
Optic' k is TxSkel a -> a -> m ()
setTweak Optic' A_Lens NoIx TxSkel [TxSkelOut]
txSkelOutsL [TxSkelOut]
newOutputs
  [TxSkelOut] -> m [TxSkelOut]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ([TxSkelOut] -> m [TxSkelOut]) -> [TxSkelOut] -> m [TxSkelOut]
forall a b. (a -> b) -> a -> b
$ [Maybe TxSkelOut] -> [TxSkelOut]
forall a. [Maybe a] -> [a]
catMaybes [Maybe TxSkelOut]
changed
  where
    go :: [TxSkelOut] -> Integer -> [(Maybe TxSkelOut, TxSkelOut)]
go [] Integer
_ = []
    go (TxSkelOut
out : [TxSkelOut]
l) Integer
n =
      case Optic' An_AffineTraversal NoIx TxSkelOut owner
-> TxSkelOut -> Maybe owner
forall k (is :: IxList) s a.
Is k An_AffineFold =>
Optic' k is s a -> s -> Maybe a
preview (forall a. OwnerConstrs a => AffineTraversal' TxSkelOut a
txSkelOutTypedOwnerAT @owner) TxSkelOut
out Maybe owner -> Maybe owner' -> Maybe owner'
forall a b. Maybe a -> Maybe b -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> TxSkelOut -> Maybe owner'
outputPred TxSkelOut
out of
        Maybe owner'
Nothing -> (Maybe TxSkelOut
forall a. Maybe a
Nothing, TxSkelOut
out) (Maybe TxSkelOut, TxSkelOut)
-> [(Maybe TxSkelOut, TxSkelOut)] -> [(Maybe TxSkelOut, TxSkelOut)]
forall a. a -> [a] -> [a]
: [TxSkelOut] -> Integer -> [(Maybe TxSkelOut, TxSkelOut)]
go [TxSkelOut]
l Integer
n
        Just owner'
newOwner | Integer -> Bool
indexPred Integer
n -> (TxSkelOut -> Maybe TxSkelOut
forall a. a -> Maybe a
Just TxSkelOut
out, TxSkelOut
out {tsoOwner = newOwner}) (Maybe TxSkelOut, TxSkelOut)
-> [(Maybe TxSkelOut, TxSkelOut)] -> [(Maybe TxSkelOut, TxSkelOut)]
forall a. a -> [a] -> [a]
: [TxSkelOut] -> Integer -> [(Maybe TxSkelOut, TxSkelOut)]
go [TxSkelOut]
l (Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1)
        Maybe owner'
_ -> (Maybe TxSkelOut
forall a. Maybe a
Nothing, TxSkelOut
out) (Maybe TxSkelOut, TxSkelOut)
-> [(Maybe TxSkelOut, TxSkelOut)] -> [(Maybe TxSkelOut, TxSkelOut)]
forall a. a -> [a] -> [a]
: [TxSkelOut] -> Integer -> [(Maybe TxSkelOut, TxSkelOut)]
go [TxSkelOut]
l (Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1)

-- | A version of 'redirectOutputTweakAll' where, instead of modifying all the
-- outputs targeted by the input predicates in the same transaction, we modify
-- one of them at a time, relying on the 'MonadPlus' instance of @m@.
redirectOutputTweakAny ::
  forall owner owner' m.
  (MonadTweak m, OwnerConstrs owner, OwnerConstrs owner') =>
  (TxSkelOut -> Maybe owner') ->
  (Integer -> Bool) ->
  m TxSkelOut
redirectOutputTweakAny :: forall owner owner' (m :: * -> *).
(MonadTweak m, OwnerConstrs owner, OwnerConstrs owner') =>
(TxSkelOut -> Maybe owner') -> (Integer -> Bool) -> m TxSkelOut
redirectOutputTweakAny TxSkelOut -> Maybe owner'
outputPred Integer -> Bool
indexPred = Optic' A_Lens NoIx TxSkel [TxSkelOut] -> m [TxSkelOut]
forall (m :: * -> *) k (is :: IxList) a.
(MonadTweak m, Is k A_Getter) =>
Optic' k is TxSkel a -> m a
viewTweak Optic' A_Lens NoIx TxSkel [TxSkelOut]
txSkelOutsL m [TxSkelOut] -> ([TxSkelOut] -> m TxSkelOut) -> m TxSkelOut
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= [TxSkelOut] -> Integer -> [TxSkelOut] -> m TxSkelOut
go [] Integer
0
  where
    go :: [TxSkelOut] -> Integer -> [TxSkelOut] -> m TxSkelOut
go [TxSkelOut]
_ Integer
_ [] = m TxSkelOut
forall a. m a
forall (m :: * -> *) a. MonadPlus m => m a
mzero
    go [TxSkelOut]
l' Integer
n (TxSkelOut
out : [TxSkelOut]
l)
      | Integer -> Bool
indexPred Integer
n =
          m TxSkelOut -> Maybe (m TxSkelOut) -> m TxSkelOut
forall a. a -> Maybe a -> a
fromMaybe
            ([TxSkelOut] -> Integer -> [TxSkelOut] -> m TxSkelOut
go ([TxSkelOut]
l' [TxSkelOut] -> [TxSkelOut] -> [TxSkelOut]
forall a. [a] -> [a] -> [a]
++ [TxSkelOut
out]) (Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1) [TxSkelOut]
l)
            ( do
                Maybe owner -> Maybe ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (Maybe owner -> Maybe ()) -> Maybe owner -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Optic' An_AffineTraversal NoIx TxSkelOut owner
-> TxSkelOut -> Maybe owner
forall k (is :: IxList) s a.
Is k An_AffineFold =>
Optic' k is s a -> s -> Maybe a
preview (forall a. OwnerConstrs a => AffineTraversal' TxSkelOut a
txSkelOutTypedOwnerAT @owner) TxSkelOut
out
                owner'
newOwner <- TxSkelOut -> Maybe owner'
outputPred TxSkelOut
out
                m TxSkelOut -> Maybe (m TxSkelOut)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (m TxSkelOut -> Maybe (m TxSkelOut))
-> m TxSkelOut -> Maybe (m TxSkelOut)
forall a b. (a -> b) -> a -> b
$
                  m TxSkelOut -> m TxSkelOut -> m TxSkelOut
forall a. m a -> m a -> m a
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
mplus
                    (Optic' A_Lens NoIx TxSkel [TxSkelOut] -> [TxSkelOut] -> m ()
forall (m :: * -> *) k (is :: IxList) a.
(MonadTweak m, Is k A_Setter) =>
Optic' k is TxSkel a -> a -> m ()
setTweak Optic' A_Lens NoIx TxSkel [TxSkelOut]
txSkelOutsL ([TxSkelOut]
l' [TxSkelOut] -> [TxSkelOut] -> [TxSkelOut]
forall a. [a] -> [a] -> [a]
++ TxSkelOut
out {tsoOwner = newOwner} TxSkelOut -> [TxSkelOut] -> [TxSkelOut]
forall a. a -> [a] -> [a]
: [TxSkelOut]
l) m () -> m TxSkelOut -> m TxSkelOut
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> TxSkelOut -> m TxSkelOut
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return TxSkelOut
out)
                    ([TxSkelOut] -> Integer -> [TxSkelOut] -> m TxSkelOut
go ([TxSkelOut]
l' [TxSkelOut] -> [TxSkelOut] -> [TxSkelOut]
forall a. [a] -> [a] -> [a]
++ [TxSkelOut
out]) (Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1) [TxSkelOut]
l)
            )
    go [TxSkelOut]
l' Integer
n (TxSkelOut
out : [TxSkelOut]
l) = [TxSkelOut] -> Integer -> [TxSkelOut] -> m TxSkelOut
go ([TxSkelOut]
l' [TxSkelOut] -> [TxSkelOut] -> [TxSkelOut]
forall a. [a] -> [a] -> [a]
++ [TxSkelOut
out]) Integer
n [TxSkelOut]
l

-- | A datum hijacking attack, simplified: This attack tries to substitute a
-- different recipient on outputs belonging to scripts, but leaves the datum as
-- it is. That is, it tests for careless uses of something like
-- 'Api.txInfoOutputs' in places where something like 'Api.getContinuingOutputs'
-- should be used. If this attack goes through, however, a "proper" datum
-- hijacking attack that modifies the datum in a way that (the relevant part of)
-- the 'Api.toBuiltinData'-translation stays the same will also work.
--
-- A 'DatumHijackingLbl' with the hash of the "thief" validator is added to the
-- labels of the 'TxSkel' using 'addLabelTweak'.
--
-- This attack returns the list of outputs it redirected, in the order in which
-- they occurred on the original transaction. If no output is redirected, this
-- attack fails.
datumHijackingAttackAll ::
  forall owner owner' m.
  (MonadTweak m, OwnerConstrs owner, OwnerConstrs owner') =>
  -- | Predicate to select outputs to steal, depending on the intended
  -- recipient, the datum, and the value.
  (TxSkelOut -> Bool) ->
  -- | The selection predicate may match more than one output. Use this
  -- predicate to restrict to the i-th of the outputs (counting from the left,
  -- starting at zero) chosen by the selection predicate with this predicate.
  (Integer -> Bool) ->
  -- | The thief
  owner' ->
  m [TxSkelOut]
datumHijackingAttackAll :: forall owner owner' (m :: * -> *).
(MonadTweak m, OwnerConstrs owner, OwnerConstrs owner') =>
(TxSkelOut -> Bool) -> (Integer -> Bool) -> owner' -> m [TxSkelOut]
datumHijackingAttackAll TxSkelOut -> Bool
change Integer -> Bool
select owner'
thief = do
  [TxSkelOut]
redirected <- forall owner owner' (m :: * -> *).
(MonadTweak m, OwnerConstrs owner, OwnerConstrs owner') =>
(TxSkelOut -> Maybe owner') -> (Integer -> Bool) -> m [TxSkelOut]
redirectOutputTweakAll @owner (\TxSkelOut
output -> if TxSkelOut -> Bool
change TxSkelOut
output then owner' -> Maybe owner'
forall a. a -> Maybe a
Just owner'
thief else Maybe owner'
forall a. Maybe a
Nothing) Integer -> Bool
select
  Bool -> m ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> m ()) -> (Bool -> Bool) -> Bool -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> Bool
not (Bool -> m ()) -> Bool -> m ()
forall a b. (a -> b) -> a -> b
$ [TxSkelOut] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TxSkelOut]
redirected
  DatumHijackingLbl -> m ()
forall (m :: * -> *) x. (MonadTweak m, LabelConstrs x) => x -> m ()
addLabelTweak (DatumHijackingLbl -> m ()) -> DatumHijackingLbl -> m ()
forall a b. (a -> b) -> a -> b
$ Credential -> DatumHijackingLbl
DatumHijackingLbl (Credential -> DatumHijackingLbl)
-> Credential -> DatumHijackingLbl
forall a b. (a -> b) -> a -> b
$ owner' -> Credential
forall a. ToCredential a => a -> Credential
Script.toCredential owner'
thief
  [TxSkelOut] -> m [TxSkelOut]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return [TxSkelOut]
redirected

-- | A version of datumHijackingAttackAll relying on the rules of
-- 'redirectOutputTweakAny'.
datumHijackingAttackAny ::
  forall owner owner' m.
  (MonadTweak m, OwnerConstrs owner, OwnerConstrs owner') =>
  -- | Predicate to select outputs to steal, depending on the intended
  -- recipient, the datum, and the value.
  (TxSkelOut -> Bool) ->
  -- | The selection predicate may match more than one output. Use this
  -- predicate to restrict to the i-th of the outputs (counting from the left,
  -- starting at zero) chosen by the selection predicate with this predicate.
  (Integer -> Bool) ->
  -- | The thief
  owner' ->
  m TxSkelOut
datumHijackingAttackAny :: forall owner owner' (m :: * -> *).
(MonadTweak m, OwnerConstrs owner, OwnerConstrs owner') =>
(TxSkelOut -> Bool) -> (Integer -> Bool) -> owner' -> m TxSkelOut
datumHijackingAttackAny TxSkelOut -> Bool
change Integer -> Bool
select owner'
thief = do
  TxSkelOut
redirected <- forall owner owner' (m :: * -> *).
(MonadTweak m, OwnerConstrs owner, OwnerConstrs owner') =>
(TxSkelOut -> Maybe owner') -> (Integer -> Bool) -> m TxSkelOut
redirectOutputTweakAny @owner (\TxSkelOut
output -> if TxSkelOut -> Bool
change TxSkelOut
output then owner' -> Maybe owner'
forall a. a -> Maybe a
Just owner'
thief else Maybe owner'
forall a. Maybe a
Nothing) Integer -> Bool
select
  DatumHijackingLbl -> m ()
forall (m :: * -> *) x. (MonadTweak m, LabelConstrs x) => x -> m ()
addLabelTweak (DatumHijackingLbl -> m ()) -> DatumHijackingLbl -> m ()
forall a b. (a -> b) -> a -> b
$ Credential -> DatumHijackingLbl
DatumHijackingLbl (Credential -> DatumHijackingLbl)
-> Credential -> DatumHijackingLbl
forall a b. (a -> b) -> a -> b
$ owner' -> Credential
forall a. ToCredential a => a -> Credential
Script.toCredential owner'
thief
  TxSkelOut -> m TxSkelOut
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return TxSkelOut
redirected

-- | The default datum hijacking attack. It tries to redirect any output for
-- which the owner is of type @owner@ and branches at each attempt.
datumHijackingAttack ::
  forall owner owner' m.
  (MonadTweak m, OwnerConstrs owner, OwnerConstrs owner') =>
  owner' ->
  m TxSkelOut
datumHijackingAttack :: forall owner owner' (m :: * -> *).
(MonadTweak m, OwnerConstrs owner, OwnerConstrs owner') =>
owner' -> m TxSkelOut
datumHijackingAttack = forall owner owner' (m :: * -> *).
(MonadTweak m, OwnerConstrs owner, OwnerConstrs owner') =>
(TxSkelOut -> Bool) -> (Integer -> Bool) -> owner' -> m TxSkelOut
datumHijackingAttackAny @owner (Bool -> TxSkelOut -> Bool
forall a b. a -> b -> a
const Bool
True) (Bool -> Integer -> Bool
forall a b. a -> b -> a
const Bool
True)

-- | A label that is added to a 'TxSkel' that has successfully been modified by
-- any of the datum hijacking attacks
newtype DatumHijackingLbl = DatumHijackingLbl Api.Credential
  deriving (Int -> DatumHijackingLbl -> ShowS
[DatumHijackingLbl] -> ShowS
DatumHijackingLbl -> String
(Int -> DatumHijackingLbl -> ShowS)
-> (DatumHijackingLbl -> String)
-> ([DatumHijackingLbl] -> ShowS)
-> Show DatumHijackingLbl
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DatumHijackingLbl -> ShowS
showsPrec :: Int -> DatumHijackingLbl -> ShowS
$cshow :: DatumHijackingLbl -> String
show :: DatumHijackingLbl -> String
$cshowList :: [DatumHijackingLbl] -> ShowS
showList :: [DatumHijackingLbl] -> ShowS
Show, DatumHijackingLbl -> DatumHijackingLbl -> Bool
(DatumHijackingLbl -> DatumHijackingLbl -> Bool)
-> (DatumHijackingLbl -> DatumHijackingLbl -> Bool)
-> Eq DatumHijackingLbl
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: DatumHijackingLbl -> DatumHijackingLbl -> Bool
== :: DatumHijackingLbl -> DatumHijackingLbl -> Bool
$c/= :: DatumHijackingLbl -> DatumHijackingLbl -> Bool
/= :: DatumHijackingLbl -> DatumHijackingLbl -> Bool
Eq, Eq DatumHijackingLbl
Eq DatumHijackingLbl =>
(DatumHijackingLbl -> DatumHijackingLbl -> Ordering)
-> (DatumHijackingLbl -> DatumHijackingLbl -> Bool)
-> (DatumHijackingLbl -> DatumHijackingLbl -> Bool)
-> (DatumHijackingLbl -> DatumHijackingLbl -> Bool)
-> (DatumHijackingLbl -> DatumHijackingLbl -> Bool)
-> (DatumHijackingLbl -> DatumHijackingLbl -> DatumHijackingLbl)
-> (DatumHijackingLbl -> DatumHijackingLbl -> DatumHijackingLbl)
-> Ord DatumHijackingLbl
DatumHijackingLbl -> DatumHijackingLbl -> Bool
DatumHijackingLbl -> DatumHijackingLbl -> Ordering
DatumHijackingLbl -> DatumHijackingLbl -> DatumHijackingLbl
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: DatumHijackingLbl -> DatumHijackingLbl -> Ordering
compare :: DatumHijackingLbl -> DatumHijackingLbl -> Ordering
$c< :: DatumHijackingLbl -> DatumHijackingLbl -> Bool
< :: DatumHijackingLbl -> DatumHijackingLbl -> Bool
$c<= :: DatumHijackingLbl -> DatumHijackingLbl -> Bool
<= :: DatumHijackingLbl -> DatumHijackingLbl -> Bool
$c> :: DatumHijackingLbl -> DatumHijackingLbl -> Bool
> :: DatumHijackingLbl -> DatumHijackingLbl -> Bool
$c>= :: DatumHijackingLbl -> DatumHijackingLbl -> Bool
>= :: DatumHijackingLbl -> DatumHijackingLbl -> Bool
$cmax :: DatumHijackingLbl -> DatumHijackingLbl -> DatumHijackingLbl
max :: DatumHijackingLbl -> DatumHijackingLbl -> DatumHijackingLbl
$cmin :: DatumHijackingLbl -> DatumHijackingLbl -> DatumHijackingLbl
min :: DatumHijackingLbl -> DatumHijackingLbl -> DatumHijackingLbl
Ord)

instance PrettyCooked DatumHijackingLbl where
  prettyCookedOpt :: PrettyCookedOpts -> DatumHijackingLbl -> DocCooked
prettyCookedOpt PrettyCookedOpts
opts (DatumHijackingLbl Credential
address) = DocCooked
"DatumHijacking" DocCooked -> DocCooked -> DocCooked
forall ann. Doc ann -> Doc ann -> Doc ann
<+> PrettyCookedOpts -> Credential -> DocCooked
forall a. PrettyCooked a => PrettyCookedOpts -> a -> DocCooked
prettyCookedOpt PrettyCookedOpts
opts Credential
address