#
# Nim Ethereum Keys (nim-eth-keys)
# Copyright (c) 2018 Status Research & Development GmbH
# Licensed under either of
# - Apache License, version 2.0, (LICENSE-APACHEv2)
# - MIT license (LICENSE-MIT)
#

## This is libsecp256k1 backend.

import secp256k1, nimcrypto/sysrand, nimcrypto/utils

const
  KeyLength* = 256 div 8
  RawSignatureSize* = KeyLength * 2 + 1
  RawPublicKeySize* = KeyLength * 2
  InvalidPrivateKey = "Invalid private key!"
  InvalidPublicKey = "Invalid public key!"
  InvalidSignature = "Invalid signature!"
  VerificationFailed = "Signature verification has been failed!"
  EmptyMessageError = "Data/Message could not be empty!"
  MessageSizeError = "Size of message to sign must be KeyLength bytes!"

type
  PublicKey* = secp256k1_pubkey
    ## Representation of public key

  PrivateKey* = object
    ## Representation of secret key
    data*: array[KeyLength, byte]

  SharedSecret* = object
    ## Representation of ECDH shared secret
    data*: array[KeyLength, byte]

  KeyPair* = object
    ## Representation of private/public keys pair
    seckey*: PrivateKey
    pubkey*: PublicKey

  Signature* = secp256k1_ecdsa_recoverable_signature
    ## Representation of signature

  Secp256k1Exception* = object of Exception
    ## Exceptions generated by `libsecp256k1`

  EthKeysContext = ref object
    context: ptr secp256k1_context
    error: string

var ekContext {.threadvar.}: EthKeysContext
  ## Thread local variable which holds current context

##
## Private procedures interface
##

proc illegalCallback(message: cstring; data: pointer) {.cdecl.} =
  let ctx = cast[EthKeysContext](data)
  ctx.error = $message

proc errorCallback(message: cstring, data: pointer) {.cdecl.} =
  let ctx = cast[EthKeysContext](data)
  ctx.error = $message

proc newEthKeysContext(): EthKeysContext =
  ## Create new `EthKeysContext`.
  result = new EthKeysContext
  let flags = cuint(SECP256K1_CONTEXT_VERIFY or SECP256K1_CONTEXT_SIGN)
  result.context = secp256k1_context_create(flags)
  secp256k1_context_set_illegal_callback(result.context, illegalCallback,
                                         cast[pointer](result))
  secp256k1_context_set_error_callback(result.context, errorCallback,
                                       cast[pointer](result))
  result.error = ""

proc getSecpContext(): ptr secp256k1_context =
  ## Get current `secp256k1_context`
  if isNil(ekContext):
    ekContext = newEthKeysContext()
  result = ekContext.context

proc getContext(): EthKeysContext =
  ## Get current `EccContext`
  if isNil(ekContext):
    ekContext = newEthKeysContext()
  result = ekContext

template raiseSecp256k1Error() =
  ## Raises `libsecp256k1` error as exception
  let mctx = getContext()
  if len(mctx.error) > 0:
    var msg = mctx.error
    mctx.error.setLen(0)
    raise newException(Secp256k1Exception, msg)

proc libsecp256k1ErrorMsg(): string =
  let mctx = getContext()
  result = mctx.error

proc setErrorMsg(m: string) =
  let mctx = getContext()
  mctx.error = m

##
## Public procedures interface
##

proc newPrivateKey*(): PrivateKey =
  ## Generates new private key.
  let ctx = getSecpContext()
  while true:
    if randomBytes(result.data) == KeyLength:
      if secp256k1_ec_seckey_verify(ctx, cast[ptr cuchar](addr result)) == 1:
        break

proc getPublicKey*(seckey: PrivateKey): PublicKey =
  ## Return public key for private key `seckey`.
  let ctx = getSecpContext()
  if secp256k1_ec_pubkey_create(ctx, addr result,
                                cast[ptr cuchar](unsafeAddr seckey)) != 1:
    raiseSecp256k1Error()

proc newKeyPair*(): KeyPair =
  ## Generates new private and public key.
  result.seckey = newPrivateKey()
  result.pubkey = result.seckey.getPublicKey()

proc initPrivateKey*(hexstr: string): PrivateKey =
  ## Create new private key from hexadecimal string representation.
  let ctx = getSecpContext()
  var o = fromHex(stripSpaces(hexstr))
  if len(o) < KeyLength:
    raise newException(EthKeysException, InvalidPrivateKey)
  copyMem(addr result, addr o[0], KeyLength)
  if secp256k1_ec_seckey_verify(ctx, cast[ptr cuchar](addr result)) != 1:
    raise newException(EthKeysException, InvalidPrivateKey)

proc initPrivateKey*(data: openarray[byte]): PrivateKey =
  ## Create new private key from binary data blob.
  let ctx = getSecpContext()
  if len(data) < KeyLength:
    raise newException(EthKeysException, InvalidPrivateKey)
  copyMem(addr result, unsafeAddr data[0], KeyLength)
  if secp256k1_ec_seckey_verify(ctx, cast[ptr cuchar](addr result)) != 1:
    raise newException(EthKeysException, InvalidPrivateKey)

proc recoverPublicKey*(data: openarray[byte],
                       pubkey: var PublicKey): EthKeysStatus =
  ## Unserialize public key from `data`.
  let ctx = getSecpContext()
  let length = len(data)
  if length < RawPublicKeySize:
    setErrorMsg(InvalidPublicKey)
    return(EthKeysStatus.Error)
  var rawkey: array[RawPublicKeySize + 1, byte]
  rawkey[0] = 0x04'u8 # mark key with UNCOMPRESSED flag
  copyMem(addr rawkey[1], unsafeAddr data[0], RawPublicKeySize)
  if secp256k1_ec_pubkey_parse(ctx, addr pubkey,
                               cast[ptr cuchar](addr rawkey),
                               RawPublicKeySize + 1) != 1:
    return(EthKeysStatus.Error)
  result = EthKeysStatus.Success

proc recoverSignature*(data: openarray[byte],
                       signature: var Signature): EthKeysStatus =
  ## Unserialize signature from `data`.
  let ctx = getSecpContext()
  let length = len(data)
  if length < RawSignatureSize:
    setErrorMsg(InvalidSignature)
    return(EthKeysStatus.Error)
  var recid = cint(data[KeyLength * 2])
  if secp256k1_ecdsa_recoverable_signature_parse_compact(ctx, addr signature,
                                           cast[ptr cuchar](unsafeAddr data[0]),
                                                         recid) != 1:
    return(EthKeysStatus.Error)
  result = EthKeysStatus.Success

proc initPublicKey*(hexstr: string): PublicKey =
  ## Create new public key from hexadecimal string representation.
  var o = fromHex(stripSpaces(hexstr))
  if len(o) < RawPublicKeySize:
    raise newException(EthKeysException, InvalidPublicKey)
  if recoverPublicKey(o, result) != EthKeysStatus.Success:
    raise newException(EthKeysException, InvalidPublicKey)

proc initPublicKey*(data: openarray[byte]): PublicKey =
  ## Create new public key from binary data blob.
  if recoverPublicKey(data, result) != EthKeysStatus.Success:
    raise newException(EthKeysException, InvalidPublicKey)

proc initSignature*(hexstr: string): Signature =
  ## Create new signature from hexadecimal string representation.
  var o = fromHex(stripSpaces(hexstr))
  if recoverSignature(o, result) != EthKeysStatus.Success:
    raise newException(EthKeysException, libsecp256k1ErrorMsg())

proc initSignature*(data: openarray[byte]): Signature =
  ## Create new signature from 'data'.
  if recoverSignature(data, result) != EthKeysStatus.Success:
    raise newException(EthKeysException, libsecp256k1ErrorMsg())

proc ecdhAgree*(seckey: PrivateKey, pubkey: PublicKey,
                secret: var SharedSecret): EthKeysStatus =
  ## Calculate ECDH shared secret.
  var res: array[KeyLength + 1, byte]
  let ctx = getSecpContext()
  if secp256k1_ecdh_raw(ctx, cast[ptr cuchar](addr res),
                        unsafeAddr pubkey,
                        cast[ptr cuchar](unsafeAddr seckey)) != 1:
    return(EthKeysStatus.Error)
  copyMem(addr secret, addr res[1], KeyLength)
  return(EthKeysStatus.Success)

proc getRaw*(pubkey: PublicKey): array[RawPublicKeySize, byte] {.noinit.} =
  ## Converts public key `pubkey` to serialized form.
  var key: array[RawPublicKeySize + 1, byte]
  var length = csize(sizeof(key))
  let ctx = getSecpContext()
  if secp256k1_ec_pubkey_serialize(ctx, cast[ptr cuchar](addr key),
                                   addr length, unsafeAddr pubkey,
                                   SECP256K1_EC_UNCOMPRESSED) != 1:
    raiseSecp256k1Error()
  assert(length == RawPublicKeySize + 1)
  assert(key[0] == 0x04'u8)
  copyMem(addr result[0], addr key[1], RawPublicKeySize)

proc toRaw*(pubkey: PublicKey, data: var openarray[byte]) =
  ## Converts public key `pubkey` to serialized form and store it in `data`.
  var key: array[RawPublicKeySize + 1, byte]
  assert(len(data) >= RawPublicKeySize)
  var length = csize(sizeof(key))
  let ctx = getSecpContext()
  if secp256k1_ec_pubkey_serialize(ctx, cast[ptr cuchar](addr key),
                                   addr length, unsafeAddr pubkey,
                                   SECP256K1_EC_UNCOMPRESSED) != 1:
    raiseSecp256k1Error()
  assert(length == RawPublicKeySize + 1)
  assert(key[0] == 0x04'u8)
  copyMem(addr data[0], addr key[1], RawPublicKeySize)

proc getRaw*(s: Signature): array[RawSignatureSize, byte] {.noinit.} =
  ## Converts signature `s` to serialized form.
  let ctx = getSecpContext()
  var recid = cint(0)
  if secp256k1_ecdsa_recoverable_signature_serialize_compact(
    ctx, cast[ptr cuchar](unsafeAddr result), addr recid, unsafeAddr s) != 1:
    raiseSecp256k1Error()
  result[64] = uint8(recid)

proc toRaw*(s: Signature, data: var openarray[byte]) =
  ## Converts signature `s` to serialized form and store it in `data`.
  let ctx = getSecpContext()
  var recid = cint(0)
  assert(len(data) >= RawSignatureSize)
  if secp256k1_ecdsa_recoverable_signature_serialize_compact(
    ctx, cast[ptr cuchar](addr data[0]), addr recid, unsafeAddr s) != 1:
    raiseSecp256k1Error()
  data[64] = uint8(recid)

proc recoverSignatureKey*(data: openarray[byte],
                          msg: openarray[byte],
                          pubkey: var PublicKey): EthKeysStatus =
  ## Perform check on digitally signed `data` using original message `msg` and
  ## recover public key to `pubkey` on success.
  let ctx = getSecpContext()
  let length = len(data)
  if len(msg) == 0 or length == 0:
    setErrorMsg(EmptyMessageError)
    return(EthKeysStatus.Error)
  if length < RawSignatureSize:
    setErrorMsg(InvalidSignature)
    return(EthKeysStatus.Error)
  var recid = cint(data[KeyLength * 2])
  var s: secp256k1_ecdsa_recoverable_signature
  if secp256k1_ecdsa_recoverable_signature_parse_compact(ctx, addr s,
                                           cast[ptr cuchar](unsafeAddr data[0]),
                                                         recid) != 1:
    return(EthKeysStatus.Error)
  if secp256k1_ecdsa_recover(ctx, addr pubkey, addr s,
                             cast[ptr cuchar](msg)) != 1:
    setErrorMsg(VerificationFailed)
    return(EthKeysStatus.Error)
  result = EthKeysStatus.Success

proc recoverSignatureKey*(signature: Signature,
                          msg: openarray[byte],
                          pubkey: var PublicKey): EthKeysStatus =
  ## Perform check of `signature` using original message `msg` and
  ## recover public key to `pubkey` on success.
  let ctx = getSecpContext()
  if len(msg) == 0:
    setErrorMsg(EmptyMessageError)
    return(EthKeysStatus.Error)
  if secp256k1_ecdsa_recover(ctx, addr pubkey, unsafeAddr signature,
                             cast[ptr cuchar](msg)) != 1:
    setErrorMsg(VerificationFailed)
    return(EthKeysStatus.Error)
  result = EthKeysStatus.Success

proc signRawMessage*(data: openarray[byte], seckey: PrivateKey,
                     signature: var Signature): EthKeysStatus =
  ## Sign message `data` of `KeyLength` size using private key `seckey` and
  ## store result into `signature`.
  let ctx = getSecpContext()
  let length = len(data)
  if length != KeyLength:
    setErrorMsg(MessageSizeError)
    return(EthKeysStatus.Error)
  if secp256k1_ecdsa_sign_recoverable(ctx, addr signature,
                                      cast[ptr cuchar](unsafeAddr data[0]),
                                      cast[ptr cuchar](unsafeAddr seckey),
                                      nil, nil) != 1:
    return(EthKeysStatus.Error)
  return(EthKeysStatus.Success)

proc shutdownLibsecp256k1() =
  if not isNil(ekContext):
    secp256k1_context_destroy(ekContext.context)
    ekContext = nil