import base64, strutils, sha256, os
import libsodium.sodium
import libsodium.sodium_sizes
const
  usage = """
DDCATool - Tool to establish and manage a DDNet CA
usage:
  ddcatool <command>

Commands:
  genKey <keyOut> [seed]
  request <csrOut> <host> <canIssue> <keyIn>
  sign <certOut> <csrIn> <issuingHost> <privIn>
  inspect <fileIn>
"""

  certVersion = 1
  certHeader  = "#BEGIN DDCERT#"
  certFooter  = "#END DDCERT#"
  certLen     = 32 + 32 + 1 + 32 + 64
  keyHeader   = "#BEGIN DDKEY#"
  keyFooter   = "#END DDKEY#"
  keyLen      = 64
  csrHeader   = "#BEGIN DDCRL#"
  csrFooter   = "#END DDCRL#"
  csrLen      = 32 + 1 + 32 + 64

type
  Cert    = tuple[hostHash: string, issuerHash: string, canIssue: bool, pubKey: string, signature: string]
  Pub     = string
  Key     = string
  KeyPair = tuple[public: Pub, private: Key]
  Csr     = tuple[hostHash: string, canIssue: bool, pubKey: string, signature: string]

proc canIssueRaw(x: bool): char =
  result = if x: '\1'
           else: '\0'

proc generateCertContent(c: Cert): string =
  result = c[0] & c[1] & canIssueRaw(c[2]) & c[3] & c[4]

proc generateCsrContent(c: Csr): string =
  result = c[0] & canIssueRaw(c[1]) & c[2] & c[3]

template createFileImpl(lines: varargs[string]) =
  var file: File
  if not file.open(f, fmWrite):
    quit("Couldn't write file.", QuitFailure) #Couldnt be assed with exceptions
  for line in lines:
    file.writeLine(line)
  file.close

proc createCertFile(f:string, c: Cert) =
  createFileImpl(certHeader, $certVersion, encode(generateCertContent(c)), certFooter)

proc createKeyFile(f: string, c: Key) = #libsodium doesn't have nice symmetric enc :(
  createFileImpl(keyHeader, encode(c), keyFooter)

proc createCsrFile(f: string, c: Csr) =
  createFileImpl(csrHeader, encode(generateCsrContent(c)), csrFooter)

proc parseCertFile(f: string): seq[Cert] =
  result = @[]

  let buf = readFile(f)

  var cert = ""
  for line in buf.splitLines:
    if line.startsWith(certHeader):
      let version = parseInt(line.split('#')[2])
      assert(version == certVersion) #Invalid version cert in file
    elif line != certHeader and line != certFooter:
      cert.add(line)
    elif line == certFooter:
      let raw = decode(cert)
      assert(raw.len == certLen)
      let
        hostHash    = raw[0..31]   #32 bytes SHA256
        issuerHash  = raw[32..63]  #32 bytes SHA256
        canIssueRaw = raw[64]      #1  byte  bool
        pubKey      = raw[65..96]  #32 bytes ED25519 Public Key
        signature   = raw[97..160] #64 bytes ED25519 Signature
        canIssue    = if canIssueRaw == '\1': true
                                        else: false

      result.add((hostHash, issuerHash, canIssue, pubKey, signature))
      cert = ""

proc parseCsrFile(f: string): Csr =
  let
    buf = readFile(f)
    count = buf.count(csrHeader)

  assert(count == 1)

  var csr = ""
  for line in buf.splitLines:
    if line != csrHeader and line != csrFooter:
      csr.add(line)
    elif line == csrFooter:
      let raw = decode(csr)
      assert(raw.len == csrLen)
      let
        hostHash    = raw[0..31]   #32 bytes SHA256
        canIssueRaw = raw[32]      #1  byte  bool
        pubKey      = raw[33..64]  #32 byte  ED25519 Public Key
        signature   = raw[65..128] #32 byte  ED25519 Signature
        canIssue    = if canIssueRaw == '\1': true
                                        else: false
      result = (hostHash, canIssue, pubKey, signature)
      return

proc parseKeyFile(f: string): Key =
  let
    buf = readFile(f)
    count = buf.count(keyHeader)

  assert(count == 1)

  var key = ""
  for line in buf.splitLines:
    if line != keyHeader and line != keyFooter:
      key.add(line)
    elif line == keyFooter:
      let raw = decode(key)
      assert(raw.len == keyLen)

      result = raw

proc getPair(key: Key): KeyPair =
  echo key.len
  result.private = key
  result.public  = crypto_sign_ed25519_sk_to_pk(key)

proc verifySig(pubKey, message, signature: string): bool =
  result = true
  try:
    crypto_sign_verify_detached(pubKey, message, signature)
  except SodiumError:
    result = false

proc sign(csr: Csr, issuerHash: string, key: Key): Cert =
  if not verifySig(csr.pubKey, csr[0] & canIssueRaw(csr[1]) & csr[2], csr.signature):
    quit("Forged CSR file!", QuitFailure)

  result.hostHash   = csr.hostHash
  result.issuerHash = issuerHash
  result.canIssue   = csr.canIssue
  result.pubKey     = csr.pubKey
  result.signature  = crypto_sign_detached(
    secret_key = key,
    message = result[0] & result[1] & canIssueRaw(result[2]) & result[3])

proc request(hostHash: string, canIssue: bool, pair: KeyPair): Csr =
  result.hostHash  = hostHash
  result.canIssue  = canIssue
  result.pubKey    = pair.public
  result.signature = crypto_sign_detached(
    secret_key = pair.private,
    message = result[0] & canIssueRaw(result[1]) & result[2])

proc genKey(keyOut: string, seed = "") =
  let key = if seed == "": crypto_sign_keypair()[1]
                     else: crypto_sign_seed_keypair(SHA256(seed))[1]

  createKeyFile(keyOut, key)

proc requestF(csrOut: string, host: string, canIssue: bool, keyIn: string) =
  let
    hostHash = SHA256(host)
    key = parseKeyFile(keyIn)
    pair = getPair(key)
    csr = request(hostHash, canIssue, pair)

  createCsrFile(csrOut, csr)

proc signCsr(certOut: string, csrIn: string, issuingHost: string, privIn: string) =
  let
    csr = parseCsrFile(csrIn)
    key = parseKeyFile(privIn)
    issuerHash = SHA256(issuingHost)
    cert = sign(csr, issuerHash, key)

  createCertFile(certOut, cert)

proc inspectCert(f: string) =
  let cert = parseCertFile(f)
  for i, c in cert:
    echo "Cert " & $i & ':'
    echo "\tHost hash  : " & toHex(c.hostHash)
    echo "\tIssuer hash: " & toHex(c.issuerHash)
    echo "\tCan issue  : " & $c.canIssue
    echo "\tPublic Key : " & toHex(c.pubKey)
    echo "\tSignature  : " & toHex(c.signature)

proc inspectCsr(f: string) =
  let csr = parseCsrFile(f)
  echo "CSR:"
  echo "\tHost hash : " & toHex(csr.hostHash)
  echo "\tCan issue : " & $csr.canIssue
  echo "\tPublic Key: " & toHex(csr.pubKey)
  echo "\tSignature : " & toHex(csr.signature)
  echo "\tValidity  : " & $verifySig(csr.pubKey, csr[0] & canIssueRaw(csr[1]) & csr[2], csr.signature)

proc inspect(file: string) =
  let buf = readFile(file)

  for line in buf.splitLines:
    case line
    of certHeader:
      inspectCert(file)
    of csrHeader:
      inspectCsr(file)
    else: discard

if paramCount() == 0:
  echo "No command"
  quit(usage, QuitFailure)

let command = paramStr(1)
case command
of "genkey":
  if paramCount() > 3:
    echo "Too many arguments for genkey."
    quit(usage, QuitFailure)

  if paramCount() == 2:
    genKey(paramStr(2))
  else:
    genKey(paramStr(2), paramStr(3))
of "request":
  if paramCount() != 5:
    echo "Argument count invalid for request."
    quit(usage, QuitFailure)

  var canIssue = false
  case paramStr(4)
  of "true", "1":
    canIssue = true
  of "false", "0":
    canIssue = false
  else:
    echo "CanIssue can be true/1 or false/0."
    quit(usage, QuitFailure)

  requestF(paramStr(2), paramStr(3), canIssue, paramStr(5))
of "sign":
  if paramCount() != 5:
    echo "Argument count invalid for sign."
    quit(usage, QuitFailure)

  signCsr(paramStr(2), paramStr(3), paramStr(4), paramStr(5))
of "inspect":
  if paramCount() != 2:
    echo "Inspect needs one argument."
    quit(usage, QuitFailure)

  inspect(paramStr(2))
else:
  echo "Invalid command " % command
  quit(usage, QuitFailure)