# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS-IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""JWT testing service API implementations in Python."""

import datetime
import io
import json

from typing import Tuple

import grpc
import tink
from tink import cleartext_keyset_handle

from tink import jwt

from google.protobuf import duration_pb2
from google.protobuf import timestamp_pb2

from protos import testing_api_pb2
from protos import testing_api_pb2_grpc


def _to_timestamp_tuple(t: datetime.datetime) -> Tuple[int, int]:
  if not t.tzinfo:
    raise ValueError('datetime must have tzinfo')
  seconds = int(t.timestamp())
  nanos = int((t.timestamp() - seconds) * 1e9)
  return (seconds, nanos)


def _from_timestamp_proto(
    timestamp: timestamp_pb2.Timestamp) -> datetime.datetime:
  t = timestamp.seconds + (timestamp.nanos / 1e9)
  return datetime.datetime.fromtimestamp(t, datetime.timezone.utc)


def _from_duration_proto(
    duration: duration_pb2.Duration) -> datetime.timedelta:
  return datetime.timedelta(seconds=duration.seconds)


def raw_jwt_from_proto(proto_raw_jwt: testing_api_pb2.JwtToken) -> jwt.RawJwt:
  """Converts a proto JwtToken into a jwt.RawJwt."""
  type_header = None
  if proto_raw_jwt.HasField('type_header'):
    type_header = proto_raw_jwt.type_header.value
  issuer = None
  if proto_raw_jwt.HasField('issuer'):
    issuer = proto_raw_jwt.issuer.value
  subject = None
  if proto_raw_jwt.HasField('subject'):
    subject = proto_raw_jwt.subject.value
  audiences = list(proto_raw_jwt.audiences)
  if not audiences:
    audiences = None
  jwt_id = None
  if proto_raw_jwt.HasField('jwt_id'):
    jwt_id = proto_raw_jwt.jwt_id.value
  custom_claims = {}
  for name, claim in proto_raw_jwt.custom_claims.items():
    if claim.HasField('null_value'):
      custom_claims[name] = None
    elif claim.HasField('number_value'):
      custom_claims[name] = claim.number_value
    elif claim.HasField('string_value'):
      custom_claims[name] = claim.string_value
    elif claim.HasField('bool_value'):
      custom_claims[name] = claim.bool_value
    elif claim.HasField('json_object_value'):
      custom_claims[name] = json.loads(claim.json_object_value)
    elif claim.HasField('json_array_value'):
      custom_claims[name] = json.loads(claim.json_array_value)
    else:
      raise ValueError('claim %s has unknown type' % name)
  expiration = None
  if proto_raw_jwt.HasField('expiration'):
    expiration = _from_timestamp_proto(proto_raw_jwt.expiration)
  not_before = None
  if proto_raw_jwt.HasField('not_before'):
    not_before = _from_timestamp_proto(proto_raw_jwt.not_before)
  issued_at = None
  if proto_raw_jwt.HasField('issued_at'):
    issued_at = _from_timestamp_proto(proto_raw_jwt.issued_at)
  without_expiration = not expiration
  return jwt.new_raw_jwt(
      type_header=type_header,
      issuer=issuer,
      subject=subject,
      audiences=audiences,
      jwt_id=jwt_id,
      expiration=expiration,
      without_expiration=without_expiration,
      not_before=not_before,
      issued_at=issued_at,
      custom_claims=custom_claims)


def verifiedjwt_to_proto(
    verified_jwt: jwt.VerifiedJwt) -> testing_api_pb2.JwtToken:
  """Converts a jwt.VerifiedJwt into a proto JwtToken."""
  token = testing_api_pb2.JwtToken()
  if verified_jwt.has_type_header():
    token.type_header.value = verified_jwt.type_header()
  if verified_jwt.has_issuer():
    token.issuer.value = verified_jwt.issuer()
  if verified_jwt.has_subject():
    token.subject.value = verified_jwt.subject()
  if verified_jwt.has_audiences():
    token.audiences.extend(verified_jwt.audiences())
  if verified_jwt.has_jwt_id():
    token.jwt_id.value = verified_jwt.jwt_id()
  if verified_jwt.has_expiration():
    seconds, nanos = _to_timestamp_tuple(verified_jwt.expiration())
    token.expiration.seconds = seconds
    token.expiration.nanos = nanos
  if verified_jwt.has_not_before():
    seconds, nanos = _to_timestamp_tuple(verified_jwt.not_before())
    token.not_before.seconds = seconds
    token.not_before.nanos = nanos
  if verified_jwt.has_issued_at():
    seconds, nanos = _to_timestamp_tuple(verified_jwt.issued_at())
    token.issued_at.seconds = seconds
    token.issued_at.nanos = nanos
  for name in verified_jwt.custom_claim_names():
    value = verified_jwt.custom_claim(name)
    if value is None:
      token.custom_claims[name].null_value = testing_api_pb2.NULL_VALUE
    elif isinstance(value, bool):
      token.custom_claims[name].bool_value = value
    elif isinstance(value, (int, float)):
      token.custom_claims[name].number_value = value
    elif isinstance(value, str):
      token.custom_claims[name].string_value = value
    elif isinstance(value, dict):
      token.custom_claims[name].json_object_value = json.dumps(value)
    elif isinstance(value, list):
      token.custom_claims[name].json_array_value = json.dumps(value)
    else:
      raise ValueError('claim %s has unknown type' % name)
  return token


def validator_from_proto(
    proto_validator: testing_api_pb2.JwtValidator) -> jwt.JwtValidator:
  """Converts a proto JwtValidator into a JwtValidator."""
  expected_type_header = None
  if proto_validator.HasField('expected_type_header'):
    expected_type_header = proto_validator.expected_type_header.value
  expected_issuer = None
  if proto_validator.HasField('expected_issuer'):
    expected_issuer = proto_validator.expected_issuer.value
  expected_audience = None
  if proto_validator.HasField('expected_audience'):
    expected_audience = proto_validator.expected_audience.value
  fixed_now = None
  if proto_validator.HasField('now'):
    fixed_now = _from_timestamp_proto(proto_validator.now)
  clock_skew = None
  if proto_validator.HasField('clock_skew'):
    clock_skew = _from_duration_proto(proto_validator.clock_skew)
  return jwt.new_validator(
      expected_type_header=expected_type_header,
      expected_issuer=expected_issuer,
      expected_audience=expected_audience,
      ignore_type_header=proto_validator.ignore_type_header,
      ignore_issuer=proto_validator.ignore_issuer,
      ignore_audiences=proto_validator.ignore_audience,
      allow_missing_expiration=proto_validator.allow_missing_expiration,
      expect_issued_in_the_past=proto_validator.expect_issued_in_the_past,
      fixed_now=fixed_now,
      clock_skew=clock_skew)


class JwtServicer(testing_api_pb2_grpc.JwtServicer):
  """A service for signing and verifying JWTs."""

  def CreateJwtMac(
      self, request: testing_api_pb2.CreationRequest,
      context: grpc.ServicerContext) -> testing_api_pb2.CreationResponse:
    """Creates a JwtMac without using it."""
    try:
      keyset_handle = cleartext_keyset_handle.read(
          tink.BinaryKeysetReader(request.annotated_keyset.serialized_keyset))
      keyset_handle.primitive(jwt.JwtMac)
      return testing_api_pb2.CreationResponse()
    except tink.TinkError as e:
      return testing_api_pb2.CreationResponse(err=str(e))

  def CreateJwtPublicKeySign(
      self, request: testing_api_pb2.CreationRequest,
      context: grpc.ServicerContext) -> testing_api_pb2.CreationResponse:
    """Creates a JwtPublicKeySign without using it."""
    try:
      keyset_handle = cleartext_keyset_handle.read(
          tink.BinaryKeysetReader(request.annotated_keyset.serialized_keyset))
      keyset_handle.primitive(jwt.JwtPublicKeySign)
      return testing_api_pb2.CreationResponse()
    except tink.TinkError as e:
      return testing_api_pb2.CreationResponse(err=str(e))

  def CreateJwtPublicKeyVerify(
      self, request: testing_api_pb2.CreationRequest,
      context: grpc.ServicerContext) -> testing_api_pb2.CreationResponse:
    """Creates a JwtPublicKeyVerify without using it."""
    try:
      keyset_handle = cleartext_keyset_handle.read(
          tink.BinaryKeysetReader(request.annotated_keyset.serialized_keyset))
      keyset_handle.primitive(jwt.JwtPublicKeyVerify)
      return testing_api_pb2.CreationResponse()
    except tink.TinkError as e:
      return testing_api_pb2.CreationResponse(err=str(e))

  def ComputeMacAndEncode(
      self, request: testing_api_pb2.JwtSignRequest,
      context: grpc.ServicerContext) -> testing_api_pb2.JwtSignResponse:
    """Computes a MACed compact JWT."""
    try:
      keyset_handle = cleartext_keyset_handle.read(
          tink.BinaryKeysetReader(request.annotated_keyset.serialized_keyset))
      p = keyset_handle.primitive(jwt.JwtMac)
      raw_jwt = raw_jwt_from_proto(request.raw_jwt)
      signed_compact_jwt = p.compute_mac_and_encode(raw_jwt)
      return testing_api_pb2.JwtSignResponse(
          signed_compact_jwt=signed_compact_jwt)
    except tink.TinkError as e:
      return testing_api_pb2.JwtSignResponse(err=str(e))

  def VerifyMacAndDecode(
      self, request: testing_api_pb2.JwtVerifyRequest,
      context: grpc.ServicerContext) -> testing_api_pb2.JwtVerifyResponse:
    """Verifies a MAC value."""
    try:
      keyset_handle = cleartext_keyset_handle.read(
          tink.BinaryKeysetReader(request.annotated_keyset.serialized_keyset))
      validator = validator_from_proto(request.validator)
      p = keyset_handle.primitive(jwt.JwtMac)
      verified_jwt = p.verify_mac_and_decode(request.signed_compact_jwt,
                                             validator)
      return testing_api_pb2.JwtVerifyResponse(
          verified_jwt=verifiedjwt_to_proto(verified_jwt))
    except tink.TinkError as e:
      return testing_api_pb2.JwtVerifyResponse(err=str(e))

  def PublicKeySignAndEncode(
      self, request: testing_api_pb2.JwtSignRequest,
      context: grpc.ServicerContext) -> testing_api_pb2.JwtSignResponse:
    """Computes a signed compact JWT token."""
    try:
      keyset_handle = cleartext_keyset_handle.read(
          tink.BinaryKeysetReader(request.annotated_keyset.serialized_keyset))
      p = keyset_handle.primitive(jwt.JwtPublicKeySign)
      raw_jwt = raw_jwt_from_proto(request.raw_jwt)
      signed_compact_jwt = p.sign_and_encode(raw_jwt)
      return testing_api_pb2.JwtSignResponse(
          signed_compact_jwt=signed_compact_jwt)
    except tink.TinkError as e:
      return testing_api_pb2.JwtSignResponse(err=str(e))

  def PublicKeyVerifyAndDecode(
      self, request: testing_api_pb2.JwtVerifyRequest,
      context: grpc.ServicerContext) -> testing_api_pb2.JwtVerifyResponse:
    """Verifies the validity of the signed compact JWT token."""
    try:
      keyset_handle = cleartext_keyset_handle.read(
          tink.BinaryKeysetReader(request.annotated_keyset.serialized_keyset))
      validator = validator_from_proto(request.validator)
      p = keyset_handle.primitive(jwt.JwtPublicKeyVerify)
      verified_jwt = p.verify_and_decode(request.signed_compact_jwt, validator)
      return testing_api_pb2.JwtVerifyResponse(
          verified_jwt=verifiedjwt_to_proto(verified_jwt))
    except tink.TinkError as e:
      return testing_api_pb2.JwtVerifyResponse(err=str(e))

  def ToJwkSet(
      self, request: testing_api_pb2.JwtToJwkSetRequest,
      context: grpc.ServicerContext) -> testing_api_pb2.JwtToJwkSetResponse:
    """Converts a Tink Keyset with JWT keys into a JWK set."""
    try:
      keyset_handle = cleartext_keyset_handle.read(
          tink.BinaryKeysetReader(request.keyset))
      jwk_set = jwt.jwk_set_from_public_keyset_handle(keyset_handle)
      return testing_api_pb2.JwtToJwkSetResponse(jwk_set=jwk_set)
    except tink.TinkError as e:
      return testing_api_pb2.JwtToJwkSetResponse(err=str(e))

  def FromJwkSet(
      self, request: testing_api_pb2.JwtFromJwkSetRequest,
      context: grpc.ServicerContext) -> testing_api_pb2.JwtFromJwkSetResponse:
    """Converts a JWK set into a Tink Keyset."""
    try:
      keyset_handle = jwt.jwk_set_to_public_keyset_handle(request.jwk_set)
      keyset = io.BytesIO()
      cleartext_keyset_handle.write(
          tink.BinaryKeysetWriter(keyset), keyset_handle)
      return testing_api_pb2.JwtFromJwkSetResponse(keyset=keyset.getvalue())
    except tink.TinkError as e:
      return testing_api_pb2.JwtFromJwkSetResponse(err=str(e))