v4.nim

import strutils
import tables
import algorithm
import httpcore
import unicode
import strformat
import times
import nimcrypto
import httpclient
import utils

const
  decodedChars = {'a'..'z', 'A'..'Z', '0'..'9', '-', '.', '_', '~', '/'}
  newLine = '\l'
  signingAlgorithm = "AWS4-HMAC-SHA256"


# StringSegment Implementation
type StringSegment = tuple[start, mid, stop: int]


proc len(segment: StringSegment): int =
  result = segment.stop - segment.start + 1


template updates(a, b: untyped) =
  # write stdout, b
  update(a, b)


proc sortSegmentsIn(segments: var seq[StringSegment], source: string) =
  segments.sort do (first, second: StringSegment) -> int:
    let
      firstSegmentStart = first.start
      firstSegmentMid = first.mid
    let
      secondSegmentStart = second.start
      secondSegmentMid = second.mid
    result = cmp(
      toOpenArray(source, firstSegmentStart, firstSegmentMid),
      toOpenArray(source, secondSegmentStart, secondSegmentMid)
    )


# ## Canonical URI Procs ## #
proc updateWithCanonicalUri(result: var sha256, uri: string) =
  const
    hexChars = "0123456789ABCDEF"
    uriDecodedChars = decodedChars + {'/'}

  var hexContainer = ['%', '\0', '\0']
  
  for index, character in uri:
    if character in uriDecodedChars:
      updates(result, character)
    else:
      let number = ord(character)
      hexContainer[1] = hexChars[number shr 4]
      hexContainer[2] = hexChars[number and 0xF]
      updates(result, hexContainer)


# ## Canonical Query String Procs ## #
proc findQueryStringSegments(result: var seq[StringSegment], qs: string) =
  ## Find the query segments in a query string
  template addSegment(segmentStart, equalsIndex, segmentEnd) =
    # The last index and the current index must contain at least one character
    let valid = (
      segmentStart < segmentEnd and
      segmentStart < equalsIndex and
      equalsIndex <= segmentEnd
    )
    if valid:
      add(result, (segmentStart, equalsIndex, segmentEnd))

  var
    lastAmpersand = -1
    lastEquals = 0
  for index, character in qs:
    case character:
    of '&':
      addSegment(lastAmpersand+1, lastEquals, index-1)
      lastAmpersand = index
    of '=':
      lastEquals = index
    else:
      discard

  addSegment(lastAmpersand+1, lastEquals, high(qs))


proc updateWithCanonicalQueryString(result: var sha256, qs: string) =
  ## ASSUMPTIONS:
  ##   - Query string keys and values are already url encoded.
  if len(qs) == 0:
    return

  var segments = newSeq[StringSegment]()

  # Find and order all the query segments
  findQueryStringSegments(segments, qs)
  sortSegmentsIn(segments, qs)

  if len(segments) <= 0:
    return

  for index, segment in segments:
    updates(result, toOpenArray(qs, segment.start, segment.stop))
    if index != high(segments):
      updates(result, '&')


# ## Header Data Procs ## #
type HeaderData = tuple[
  data: string,
  segments: seq[StringSegment]
]

proc addLoweredString(result: var string, name: string) =
  for r in runes(name):
    let
      loweredRune = toLower(r)
      position = len(result)
    fastToUTF8Copy(loweredRune, result, position, false)


proc addCondensedString(result: var string, source: string) =
  var
    startPos = 0
    endPos = 0

  # Find the start index
  for index in countUp(0, high(source)):
    let character = source[index]
    if character != ' ':
      startPos = index
      break

  # Find the stop index
  for index in countDown(high(source), 0):
    let character = source[index]
    if character != ' ':
      endPos = index
      break

  # Handle cases where the string is entirely whitespace
  if startPos >= endPos:
    return

  # Condense spaces
  var seenSpace = false

  template addSpace =
    if seenSpace:
      seenSpace = false
      add(result, ' ')

  for character in source:
    case character
    of ' ':
      if seenSpace == false:
        seenSpace = true
    else:
      addSpace()
      add(result, character)

  addSpace()


proc createSimplifiedHeaderData(headers: HttpHeaders): HeaderData =
  result.data = ""
  result.segments = @[]

  # Add the canonical headers
  for key, values in headers.table:
    var segment: StringSegment
    segment.start = len(result.data)

    # Add the lowered version of the header
    addLoweredString(result.data, key)
    add(result.data, ':')
    segment.mid = high(result.data)

    # Add a trimmed, comma-separated string
    for value in values:
      addCondensedString(result.data, value)
      add(result.data, ',')

    # Add a newline
    setLen(result.data, high(result.data))
    segment.stop = high(result.data)
    add(result.data, '\l')

    # Add the segment
    add(result.segments, segment)

  # Order the segments
  sortSegmentsIn(result.segments, result.data)


proc updateWithCanonicalHeaders(result: var sha256, headerData: HeaderData) =
  for segment in headerData.segments:
    updates(
      result,
      toOpenArray(headerData.data, segment.start, segment.stop)
    )
    updates(result, newLine)


proc updateWithHeaderList(result: var sha256, headerData: HeaderData) =
  for index, segment in headerData.segments:
    updates(
      result,
      toOpenArray(headerData.data, segment.start, segment.mid-1)
    )
    if index != high(headerData.segments):
      updates(result, ';')


proc addHeaderList(result: var string, headerData: HeaderData) =
  for index, segment in headerData.segments:
    add(
      result,
      toOpenArray(headerData.data, segment.start, segment.mid-1)
    )
    if index != high(headerData.segments):
      add(result, ';')



# ## Hash Payload procs ## #
proc updateWithPayloadHash(result: var sha256, payload: string) =
  let payloadDigest = digest(sha256, payload)
  for dataByte in payloadDigest.data:
    updates(result, hexChar(dataByte, lowercase=true))


# ## Primary Signing Procs ## #
proc createScopeString(region, service, date: string): string =
  result = fmt("{date}/{region}/{service}/aws4_request")


proc createCanonicalRequestHash*(
    httpMethod:  string,
    uri:         string,
    queryString: string,
    payload:     string,
    headerData:  HeaderData): Sha256Digest =
  var shaContext: sha256
  init(shaContext)

  updates(shaContext, httpMethod)
  updates(shaContext, newLine)
  updateWithCanonicalUri(shaContext, uri)
  updates(shaContext, newLine)
  updateWithCanonicalQueryString(shaContext, queryString)
  updates(shaContext, newLine)
  updateWithCanonicalHeaders(shaContext, headerData)
  updates(shaContext, newLine)
  updateWithHeaderList(shaContext, headerData)
  updates(shaContext, newLine)
  updateWithPayloadHash(shaContext, payload)

  result = finish(shaContext)


proc createSignature*(
    requestDateTimeString: string,
    requestScope:          string,
    canonicalRequestHash:  Sha256Digest,
    signingKey:            Sha256Digest): string =

  var signingHmac: HMAC[sha256]
  init(signingHmac, signingKey.data)

  updates(signingHmac, signingAlgorithm)
  updates(signingHmac, newLine)
  updates(signingHmac, requestDateTimeString)
  updates(signingHmac, newLine)
  updates(signingHmac, requestScope)
  updates(signingHmac, newLine)

  for dataByte in canonicalRequestHash.data:
    updates(signingHmac, hexChar(dataByte, lowercase=true))

  result = `$`(finish(signingHmac), lowercase=true)


proc createSigningKey(
    secretKey:         string,
    requestDateString: string,
    region:            string,
    service:           string): Sha256Digest =
  ## Creates the signing key used to generate the authorization signature.
  ## TODO Investigate in-place hmac generation. 
  let
    dateKey    = hmac(sha256, "AWS4" & secretKey, requestDateString)
    regionKey  = hmac(sha256, dateKey.data, region)
    serviceKey = hmac(sha256, regionKey.data, service)
  
  result = hmac(sha256, serviceKey.data, "aws4_request")


proc createAwsAuthorization*(
    httpMethod:  string,
    uri:         string,
    queryString: string,
    payload:     string,
    headers:     HttpHeaders,
    region:      string,
    service:     string,
    secretKey:   string,
    accessKey:   string): string =
  ## Creates the authorization string required by requests to the AWS API.
  ## The produced string can either be used in the "X-Amz-Signature" query
  ## parameter or the "Authorization" header.
  ## NOTE:
  ##   This procedure makes the following assumptions about the parameters:
  ##     - The names and values of the query string parameters have been
  ##       URL encoded.
  ##     - The query string does *not* start with an '?'.
  ##     - The URI has been formatted for use by the desired service.
  ##       This usually means normalization.
  ## TODO: Assume input header values have been stripped & trimmed?
  let
    requestDateTimeString = headers["x-amz-date", 0]
    requestDateString = requestDateTimeString[0..7]

    headerData = createSimplifiedHeaderData(headers)
    scopeString = createScopeString(region, service, requestDateString)

    canonicalRequestHash = createCanonicalRequestHash(
      httpMethod  = httpMethod,
      uri         = uri,
      queryString = queryString,
      payload     = payload,
      headerData  = headerData,
    )
    signingKey = createSigningKey(
      secretKey         = secretKey,
      requestDateString = requestDateString,
      region            = region,
      service           = service
    )
    signature = createSignature(
      requestDateTimeString = requestDateTimeString,
      requestScope          = scopeString,
      canonicalRequestHash  = canonicalRequestHash,
      signingKey            = signingKey
    )

  # echo "canonicalRequest",     "\n", canonicalRequest,     "\n"
  # echo "canonicalRequestHash", "\n", canonicalRequestHash, "\n"
  # echo "stringToSign",         "\n", stringToSign,         "\n"
  # echo "signingKey",           "\n", signingKey,           "\n"
  # echo "signature",            "\n", signature,            "\n"
  # echo "scopeString",          "\n", scopeString,          "\n"

  result = fmt(
    "{signingAlgorithm} " &
    "Credential={accessKey}/{scopeString}, "
  )
  result.add("SignedHeaders=")
  addHeaderList(result, headerData)
  result.add(", Signature=")
  result.add(signature)