content: module Main where
  
  import           Control.Monad              (void)
  import           Crypto.Hash.BLAKE2.BLAKE2s (hash)
  import           Data.ByteArray             (ScrubbedBytes, convert)
  import           Data.ByteString            (ByteString, replicate, take, drop)
  import qualified Data.ByteString.Base16     as B16
  import qualified Data.ByteString.Base64     as B64
  import           Data.Maybe                 (fromMaybe)
  import           Data.Monoid                ((<>))
  import qualified Data.Serialize             as S
  import           Network.Socket
  import qualified Network.Socket.ByteString  as NBS
  import           Prelude hiding             (replicate, take, drop)
  
  import Crypto.Noise
  import Crypto.Noise.Cipher
  import Crypto.Noise.Cipher.ChaChaPoly1305
  import Crypto.Noise.DH
  import Crypto.Noise.DH.Curve25519
  import Crypto.Noise.HandshakePatterns (noiseIKpsk2)
  import Crypto.Noise.Hash hiding (hash)
  import Crypto.Noise.Hash.BLAKE2s
  
  import Data.Time.TAI64
  
  sampleICMPRequest :: ByteString
  sampleICMPRequest = fst . B16.decode $
    "450000250000000014018f5b0abd81020abd810108001bfa039901b6576972654775617264"
  
  validateICMPResponse :: ByteString
                       -> Bool
  validateICMPResponse r =
    -- Strip off part of IPv4 header because this is only a demo.
    drop 12 sample == drop 12 r
    where
      sample = fst . B16.decode $ "45000025e3030000400180570abd81010abd8102000023fa039901b65769726547756172640000000000000000000000"
  
  unsafeMessage :: (Cipher c, DH d, Hash h)
                => Bool
                -> Maybe ScrubbedBytes
                -> ScrubbedBytes
                -> NoiseState c d h
                -> (ScrubbedBytes, NoiseState c d h)
  unsafeMessage write mpsk msg ns = case operation msg ns of
    NoiseResultMessage ct ns' -> (ct, ns')
  
    NoiseResultNeedPSK ns' -> case mpsk of
      Nothing -> error "psk required but not provided"
      Just k  -> case operation k ns' of
        NoiseResultMessage ct ns'' -> (ct, ns'')
        _ -> error "something terrible happened"
  
    _ -> error "something terrible happened"
    where
      operation = if write then writeMessage else readMessage
  
  main :: IO ()
  main = do
    let ip           = "demo.wireguard.com"
        port         = "12913"
        myKeyB64     = "WAmgVYXkbT2bCtdcDwolI88/iVi/aV3/PHcUBTQSYmo=" -- private key
        serverKeyB64 = "qRCwZSKInrMAq5sepfCdaCsRJaoLe5jhtzfiw7CjbwM=" -- public key
        pskB64       = "FpCyhws9cxwWoV4xELtfJvjJN+zQVRPISllRWgeopVE="
  
    addrInfo <- head <$> getAddrInfo Nothing (Just ip) (Just port)
    sock     <- socket (addrFamily addrInfo) Datagram defaultProtocol
  
    let addr        = addrAddress addrInfo
        myStaticKey = fromMaybe (error "invalid private key")
                      . dhBytesToPair
                      . convert
                      . either (error "error Base64 decoding my private key") id
                      . B64.decode
                      $ myKeyB64 :: KeyPair Curve25519
  
        serverKey   = fromMaybe (error "invalid public key")
                      . dhBytesToPub
                      . convert
                      . either (error "error Base64 decoding server public key") id
                      . B64.decode
                      $ serverKeyB64 :: PublicKey Curve25519
  
        psk         = convert
                      . either (error "error decoding PSK") id
                      . B64.decode
                      $ pskB64 :: ScrubbedBytes
  
    myEphemeralKey <- dhGenKey
  
    let dho  = defaultHandshakeOpts InitiatorRole "WireGuard v1 zx2c4 Jason@zx2c4.com"
        opts = setLocalEphemeral (Just myEphemeralKey)
               . setLocalStatic  (Just myStaticKey)
               . setRemoteStatic (Just serverKey)
               $ dho
        ns0  = noiseState opts noiseIKpsk2 :: NoiseState ChaChaPoly1305 Curve25519 BLAKE2s
  
    tai64n <- convert . S.encode <$> getCurrentTAI64N
  
    -- Handshake: Initiator to responder -----------------------------------------
  
    let (msg0, ns1) = unsafeMessage True Nothing tai64n ns0
        macKey      = hash 32 mempty $ "mac1----" `mappend` (convert . dhPubToBytes) serverKey
        initiation  = "\x01\x00\x00\x00\x1c\x00\x00\x00" <> convert msg0 -- sender index = 28 to match other examples
        mac1        = hash 16 macKey initiation
  
    void $ NBS.sendTo sock (initiation <> mac1 <> replicate 16 0) addr
  
    -- Handshake: Responder to initiator -----------------------------------------
  
    (response0, _) <- NBS.recvFrom sock 1024
  
    let theirIndex  = take 4  . drop 4  $ response0
        (_, ns2)    = unsafeMessage False (Just psk) (convert . take 48 . drop 12 $ response0) ns1
  
    -- ICMP: Initiator to responder ----------------------------------------------
  
    let (msg1, ns3) = unsafeMessage True Nothing (convert sampleICMPRequest) ns2
        icmp        = "\x04\x00\x00\x00" <> theirIndex <> replicate 8 0 <> convert msg1
  
    void $ NBS.sendTo sock icmp addr
  
    -- ICMP: Responder to initiator ----------------------------------------------
  
    (response1, _) <- NBS.recvFrom sock 1024
  
    let (icmpPayload, ns4) = unsafeMessage False Nothing (convert . drop 16 $ response1) ns3
  
    -- KeepAlive: Initiator to responder -----------------------------------------
  
    if validateICMPResponse . convert $ icmpPayload
      then do
        let (msg2, _) = unsafeMessage True Nothing mempty ns4
            keepAlive = "\x04\x00\x00\x00" <> theirIndex <> "\x01" <> replicate 7 0 <> convert msg2
  
        void $ NBS.sendTo sock keepAlive addr
  
      else error "unexpected ICMP response from server!"
  