Nip04.hs

{-# LANGUAGE OverloadedStrings #-}

{-|
Module      : Nostr.Nip04
Description : NIP-04 Encrypted Direct Message (Legacy)
Copyright   : (c) Emre YILMAZ, 2026
License     : MIT
Maintainer  : z@emre.xyz

This module provides NIP-04 encryption and decryption.
-}

module Nostr.Nip04
  ( encryptNip04
  , decryptNip04
  , sharedSecret
  ) where

import Crypto.Cipher.AES (AES256)
import Crypto.Cipher.Types (BlockCipher(..), Cipher(..), IV, makeIV)
import Crypto.Error (CryptoFailable(..))
import qualified Crypto.Curve.Secp256k1 as Secp
import Data.ByteArray (ByteArray, ByteArrayAccess)
import qualified Data.ByteArray as BA
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Base64 as B64
import Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import qualified System.Entropy as E

import Nostr.Crypto (SecKey, hexToBytes, widerToInteger)
import Nostr.Event (PubKey(..))

-- | Compute shared secret using ECDH
sharedSecret :: SecKey -> PubKey -> Maybe ByteString
sharedSecret secKey pubKey = do
  -- Parse the x-only public key from hex. We prefix with 0x02 to parse as compressed point.
  pubBytes <- either (const Nothing) Just (hexToBytes $ unPubKey pubKey)
  let compressedPubKey = BS.cons 0x02 pubBytes
  pubPoint <- Secp.parse_point compressedPubKey
  
  -- just returning the X coordinate, we might need to use `mul` directly.
  -- Let's try `ecdh` first.
  Secp.ecdh pubPoint (widerToInteger secKey)

-- | PKCS7 padding
padPKCS7 :: Int -> ByteString -> ByteString
padPKCS7 blockSize bs =
  let len = BS.length bs
      padLen = blockSize - (len `mod` blockSize)
      padding = BS.replicate padLen (fromIntegral padLen)
  in BS.append bs padding

-- | PKCS7 unpadding
unpadPKCS7 :: ByteString -> Maybe ByteString
unpadPKCS7 bs =
  let len = BS.length bs
      padLen = fromIntegral $ BS.last bs
  in if padLen > 0 && padLen <= 16 && len >= padLen
     then Just $ BS.take (len - padLen) bs
     else Nothing

-- | Encrypt text message
encryptNip04 :: SecKey -> PubKey -> Text -> IO (Maybe Text)
encryptNip04 secKey pubKey msg = do
  case sharedSecret secKey pubKey of
    Nothing -> return Nothing
    Just secret -> do
      ivBytes <- E.getEntropy 16
      case initCipher secret of
        CryptoFailed _ -> return Nothing
        CryptoPassed cipher ->
          case makeIV ivBytes :: Maybe (IV AES256) of
            Nothing -> return Nothing
            Just iv -> do
              let paddedMsg = padPKCS7 16 (TE.encodeUtf8 msg)
              let ciphertext = cbcEncrypt cipher iv paddedMsg
              let b64Cipher = TE.decodeUtf8 $ B64.encode ciphertext
              let b64Iv = TE.decodeUtf8 $ B64.encode ivBytes
              return $ Just $ b64Cipher <> "?iv=" <> b64Iv

-- | Decrypt text message
decryptNip04 :: SecKey -> PubKey -> Text -> Maybe Text
decryptNip04 secKey pubKey encryptedMsg = do
  let parts = T.splitOn "?iv=" encryptedMsg
  case parts of
    [b64Cipher, b64Iv] -> do
      secret <- sharedSecret secKey pubKey
      cipherBytes <- either (const Nothing) Just (B64.decode $ TE.encodeUtf8 b64Cipher)
      ivBytes <- either (const Nothing) Just (B64.decode $ TE.encodeUtf8 b64Iv)
      
      cipher <- case initCipher secret of
        CryptoPassed c -> Just c
        CryptoFailed _ -> Nothing
        
      iv <- makeIV ivBytes :: Maybe (IV AES256)
      
      let decryptedPadded = cbcDecrypt cipher iv cipherBytes
      decryptedBytes <- unpadPKCS7 decryptedPadded
      Just $ TE.decodeUtf8 decryptedBytes
      
    _ -> Nothing

-- | Initialize AES256 cipher
initCipher :: ByteString -> CryptoFailable AES256
initCipher = cipherInit