-- SPDX-License-Identifier: LicenseRef-Sovereign-Leviathan-SL-AGPL3-001
-- Copyright (c) 2026 Ahmad Ali Parr

module CloneGate

import Data.Vect
import Data.Bits
import Decidable.Equality

%default total

------------------------------------------------------------------------
-- 1. Bytes, AES-style state, fingerprints
------------------------------------------------------------------------

public export
Byte : Type
Byte = Bits8

public export
AESState : Type
AESState = Vect 4 (Vect 4 Byte)

public export
Fingerprint : Type
Fingerprint = Vect 16 Byte

-- | Toy S-box. Real AES uses a fixed 256-byte permutation; here we
-- | just add 1 so the transformation is total and testable.
export
sbox : Byte -> Byte
sbox b = b + 1

export
subBytes : AESState -> AESState
subBytes = map (map sbox)

-- | XOR the key into the state, cell by cell.
export
addRoundKey : AESState -> AESState -> AESState
addRoundKey = zipWith (zipWith xor)

-- | One round of the toy cipher.
export
aesRound : AESState -> AESState -> AESState
aesRound st k = addRoundKey (subBytes st) k

-- | Toy "decryption": the inverse XOR.
export
decryptChallenge : AESState -> AESState -> AESState
decryptChallenge k ct = addRoundKey ct k

-- | Flatten the 4x4 state into a 16-byte fingerprint.
-- | The type-checker reduces 4 + (4 + (4 + 4)) to 16 by computation.
export
flatten : AESState -> Fingerprint
flatten (r1 :: r2 :: r3 :: r4 :: Nil) = r1 ++ r2 ++ r3 ++ r4

-- | The cryptographic gate: does the decrypted challenge match
-- | the fingerprint bound to the license?
export
verifyCrypto : Fingerprint -> AESState -> AESState -> Bool
verifyCrypto fp key challenge =
  flatten (decryptChallenge key challenge) == fp

------------------------------------------------------------------------
-- 2. License status and the type-level proof
------------------------------------------------------------------------

public export
data LicenseStatus = Trial | Expired | Valid

-- | The single-constructor proof type.  It is inhabited ONLY for
-- | `Valid`, which is what makes the whole scheme type-level: the
-- | compiler can only find a `LicenseProof Valid` if the license
-- | really is valid.
export
data LicenseProof : LicenseStatus -> Type where
  ValidProof : LicenseProof Valid

public export
data LicenseToken : Type where
  MkToken : (status : LicenseStatus)
         -> (fp     : Fingerprint)
         -> LicenseToken

-- Decidable equality on LicenseStatus (illustrative, and needed by
-- some downstream consumers).
noTrialValid : Trial = Valid -> Void
noTrialValid Refl impossible

noExpiredValid : Expired = Valid -> Void
noExpiredValid Refl impossible

export
decEqStatus : (a, b : LicenseStatus) -> Dec (a = b)
decEqStatus Trial   Trial   = Yes Refl
decEqStatus Trial   Expired = No (\case Refl impossible)
decEqStatus Trial   Valid   = No noTrialValid
decEqStatus Expired Trial   = No (\case Refl impossible)
decEqStatus Expired Expired = Yes Refl
decEqStatus Expired Valid   = No noExpiredValid
decEqStatus Valid   Trial   = No (negEqSym noTrialValid)
decEqStatus Valid   Expired = No (negEqSym noExpiredValid)
decEqStatus Valid   Valid   = Yes Refl

export
DecEq LicenseStatus where
  decEq = decEqStatus

------------------------------------------------------------------------
-- 3. The license decider -- the ONLY place ValidProof is constructed
------------------------------------------------------------------------

-- | Takes a token, a key, and a challenge response.  Returns a
-- | `Just proof` if and only if (a) the status is `Valid` and (b)
-- | the AES gate verifies the challenge against the token's
-- | fingerprint.
export
decideLicense : LicenseToken
             -> AESState
             -> AESState
             -> Maybe (LicenseProof Valid)
decideLicense (MkToken Trial   _)  _   _         = Nothing
decideLicense (MkToken Expired _)  _   _         = Nothing
decideLicense (MkToken Valid   fp) key challenge =
  case verifyCrypto fp key challenge of
    True  => Just ValidProof
    False => Nothing

------------------------------------------------------------------------
-- 4. The state-indexed, linear clone gate
------------------------------------------------------------------------

public export
data GateState = Locked | Unlocked | Spent

-- | A clone gate is indexed by its state.  External code cannot
-- | fabricate a `CloneGate Unlocked` without going through
-- | `unlockGate`.
export
data CloneGate : GateState -> Type where
  LockedGate   : (gateId : Int) -> CloneGate Locked
  UnlockedGate : (gateId : Int) -> CloneGate Unlocked
  SpentGate    : (gateId : Int) -> CloneGate Spent

-- | Unlocking consumes the locked gate linearly and requires the
-- | erased `LicenseProof Valid` witness.  The `0` multiplicity means
-- | the proof has no runtime representation.
export
unlockGate : (0 _ : LicenseProof Valid)
          -> (1 g : CloneGate Locked)
          -> CloneGate Unlocked
unlockGate _ (LockedGate i) = UnlockedGate i

-- | Cloning consumes the unlocked gate linearly and yields a payload
-- | plus a spent gate, which must then be discarded.
export
cloneGate : (1 g : CloneGate Unlocked)
         -> (Vect 4 Byte, CloneGate Spent)
cloneGate (UnlockedGate i) = (replicate 4 0, SpentGate i)

-- | The spent gate must be explicitly discarded, closing the linear
-- | protocol.
export
discardGate : (1 g : CloneGate Spent) -> ()
discardGate (SpentGate _) = ()

------------------------------------------------------------------------
-- 5. The main gate flow
------------------------------------------------------------------------

export
gateFlow : LicenseToken
        -> AESState
        -> AESState
        -> Maybe (Vect 4 Byte)
gateFlow token key challenge =
  case decideLicense token key challenge of
    Nothing    => Nothing
    Just prf =>
      let gate             = LockedGate 0
          unlocked         = unlockGate prf gate
          (payload, spent) = cloneGate unlocked
      in case discardGate spent of
           () => Just payload

------------------------------------------------------------------------
-- 6. A tiny smoke test
------------------------------------------------------------------------

export
dummyKey : AESState
dummyKey = replicate 4 (replicate 4 0)

export
dummyChallenge : AESState
dummyChallenge = replicate 4 (replicate 4 0)

export
dummyFP : Fingerprint
dummyFP = flatten (decryptChallenge dummyKey dummyChallenge)

export
validToken : LicenseToken
validToken = MkToken Valid dummyFP

export
expiredToken : LicenseToken
expiredToken = MkToken Expired dummyFP

export
testValid : Maybe (Vect 4 Byte)
testValid = gateFlow validToken dummyKey dummyChallenge

export
testExpired : Maybe (Vect 4 Byte)
testExpired = gateFlow expiredToken dummyKey dummyChallenge