import logging

from datetime import datetime

import jwt

from flask import Blueprint, jsonify, abort, request, make_response
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicNumbers
from cryptography.hazmat.primitives.asymmetric.ec import EllipticCurvePublicNumbers
from cryptography.hazmat.backends import default_backend
from jwkest.jwk import keyrep, RSAKey, ECKey

import data.model
import data.model.service_keys

from app import app
from auth.registry_jwt_auth import TOKEN_REGEX
from util.security import strictjwt


logger = logging.getLogger(__name__)
key_server = Blueprint('key_server', __name__)

JWT_HEADER_NAME = 'Authorization'
JWT_AUDIENCE = app.config['PREFERRED_URL_SCHEME'] + '://' + app.config['SERVER_HOSTNAME']


def _validate_jwk(jwk, kid):
  if 'kty' not in jwk:
    abort(400)

  if 'kid' not in jwk or jwk['kid'] != kid:
    abort(400)

  if jwk['kty'] == 'EC':
    if 'x' not in jwk or 'y' not in jwk:
      abort(400)
  elif jwk['kty'] == 'RSA':
    if 'e' not in jwk or 'n' not in jwk:
      abort(400)
  else:
    abort(400)


def _jwk_dict_to_public_key(jwk):
  jwkest_key = keyrep(jwk)
  if isinstance(jwkest_key, RSAKey):
    pycrypto_key = jwkest_key.key
    return RSAPublicNumbers(e=pycrypto_key.e, n=pycrypto_key.n).public_key(default_backend())
  elif isinstance(jwkest_key, ECKey):
    x, y = jwkest_key.get_key()
    return EllipticCurvePublicNumbers(x, y, jwkest_key.curve).public_key(default_backend())


def _validate_jwt(encoded_jwt, jwk, service):
  public_key = _jwk_dict_to_public_key(jwk)

  try:
    strictjwt.decode(encoded_jwt, public_key, algorithms=['RS256'],
                     audience=JWT_AUDIENCE, issuer=service)
  except strictjwt.InvalidTokenError:
    abort(400)


def _signer_kid(encoded_jwt):
  decoded_jwt = jwt.decode(encoded_jwt, verify=False)
  logger.debug(decoded_jwt)
  return decoded_jwt.get('signer_kid', None)


def _signer_key(service, signer_kid):
  try:
    return data.model.service_keys.get_service_key(signer_kid, service=service)
  except data.model.ServiceKeyDoesNotExist:
    abort(403)


@key_server.route('/services/<service>/keys', methods=['GET'])
def list_service_keys(service):
  keys = data.model.service_keys.list_service_keys(service)
  return jsonify({'keys': [key.jwk for key in keys]})


@key_server.route('/services/<service>/keys/<kid>', methods=['GET'])
def get_service_key(service, kid):
  try:
    key = data.model.service_keys.get_service_key(kid)
  except data.model.ServiceKeyDoesNotExist:
    abort(404)

  if key.approval is None:
    abort(409)

  return jsonify(key.jwk)


@key_server.route('/services/<service>/keys/<kid>', methods=['PUT'])
def put_service_keys(service, kid):
  expiration_date = request.args.get('expiration', None)
  if expiration_date:
    try:
      expiration_date = datetime.utcfromtimestamp(float(expiration_date))
    except ValueError:
      abort(400)

  try:
    jwk = request.get_json()
  except ValueError:
    abort(400)

  logger.debug(jwk)

  jwt_header = request.headers.get(JWT_HEADER_NAME, '')
  match = TOKEN_REGEX.match(jwt_header)
  if match is None:
    abort(400)
  encoded_jwt = match.group(1)

  _validate_jwk(jwk, kid)

  metadata = {'ip': request.remote_addr}
  signer_kid = _signer_kid(encoded_jwt)

  if kid == signer_kid or signer_kid is None:
    # The key is self-signed. Create a new instance and await approval.
    _validate_jwt(encoded_jwt, jwk, service)
    data.model.service_keys.create_service_key('', kid, service, jwk, metadata, expiration_date)
    return make_response('', 202)

  metadata.update({'created_by': 'Key Rotation'})
  signer_key = _signer_key(service, signer_kid)
  signer_jwk = signer_key.jwk
  if signer_key.service != service:
    abort(403)

  _validate_jwt(encoded_jwt, signer_jwk, service)

  try:
    data.model.service_keys.replace_service_key(signer_key.kid, kid, jwk, metadata, expiration_date)
  except data.model.ServiceKeyDoesNotExist:
    abort(404)

  return make_response('', 200)


@key_server.route('/services/<service>/keys/<kid>', methods=['DELETE'])
def delete_service_key(service, kid):
  jwt_header = request.headers.get(JWT_HEADER_NAME, '')
  match = TOKEN_REGEX.match(jwt_header)
  if match is None:
    abort(400)
  encoded_jwt = match.group(1)

  signer_kid = _signer_kid(encoded_jwt)
  signer_key = _signer_key(service, signer_kid)

  self_signed = kid == signer_kid or signer_kid == ''
  approved_key_for_service = signer_key.approval is not None

  if self_signed or approved_key_for_service:
    _validate_jwt(encoded_jwt, signer_key.jwk, service)

    try:
      data.model.service_keys.delete_service_key(service, kid)
    except data.model.ServiceKeyDoesNotExist:
      abort(404)

    return make_response('', 200)

  abort(403)