import logging
import re

from jsonschema import validate, ValidationError
from functools import wraps
from flask import request, url_for
from flask.ext.principal import identity_changed, Identity
from cryptography.x509 import load_pem_x509_certificate
from cryptography.hazmat.backends import default_backend
from cachetools import lru_cache

from app import app, get_app_url
from .auth_context import set_grant_context, get_grant_context
from .permissions import repository_read_grant, repository_write_grant
from util.names import parse_namespace_repository
from util.http import abort
from util.security import strictjwt
from data import model


logger = logging.getLogger(__name__)


TOKEN_REGEX = re.compile(r'^Bearer (([a-zA-Z0-9+/]+\.)+[a-zA-Z0-9+-_/]+)$')
ANONYMOUS_SUB = '(anonymous)'
CONTEXT_KINDS = ['user', 'token', 'oauth']

ACCESS_SCHEMA = {
  'type': 'array',
  'description': 'List of access granted to the subject',
  'items': {
    'type': 'object',
    'required': [
      'type',
      'name',
      'actions',
    ],
    'properties': {
      'type': {
        'type': 'string',
        'description': 'We only allow repository permissions',
        'enum': [
          'repository',
        ],
      },
      'name': {
        'type': 'string',
        'description': 'The name of the repository for which we are receiving access'
      },
      'actions': {
        'type': 'array',
        'description': 'List of specific verbs which can be performed against repository',
        'items': {
          'type': 'string',
          'enum': [
            'push',
            'pull',
          ],
        },
      },
    },
  },
}


class InvalidJWTException(Exception):
  pass


class GrantedEntity(object):
  def __init__(self, user=None, token=None, oauth=None):
    self.user = user
    self.token = token
    self.oauth = oauth


def get_granted_entity():
  """ Returns the entity granted in the current context, if any. Returns the GrantedEntity or None
      if none.
  """
  context = get_grant_context()
  if not context:
    return None

  kind = context.get('kind', 'anonymous')

  if not kind in CONTEXT_KINDS:
    return None

  if kind == 'user':
    user = model.user.get_user(context.get('user', ''))
    if not user:
      return None

    return GrantedEntity(user=user)

  if kind == 'token':
    token = model.token.load_token_data(context.get('token'))
    if not token:
      return None

    return GrantedEntity(token=token)

  if kind == 'oauth':
    user = model.user.get_user(context.get('user', ''))
    if not user:
      return None

    oauthtoken = model.oauth.lookup_access_token_for_user(user, context.get('oauth', ''))
    if not oauthtoken:
      return None

    return GrantedEntity(oauth=oauthtoken, user=user)

  return None


def get_granted_username():
  """ Returns the username inside the grant, if any. """
  granted = get_granted_entity()
  if not granted or not granted.user:
    return None

  return granted.user.username


def build_context_and_subject(user, token, oauthtoken):
  """ Builds the custom context field for the JWT signed token and returns it,
      along with the subject for the JWT signed token. """
  if oauthtoken:
    context = {
      'kind': 'oauth',
      'user': user.username,
      'oauth': oauthtoken.uuid,
    }

    return (context, user.username)

  if user:
    context = {
      'kind': 'user',
      'user': user.username,
    }
    return (context, user.username)

  if token:
    context = {
      'kind': 'token',
      'token': token.code,
    }
    return (context, None)

  context = {
    'kind': 'anonymous',
  }
  return (context, ANONYMOUS_SUB)


def get_auth_headers():
  """ Returns a dictionary of headers for auth responses. """
  headers = {}
  realm_auth_path = url_for('v2.generate_registry_jwt')
  authenticate = 'Bearer realm="{0}{1}",service="{2}"'.format(get_app_url(),
                                                              realm_auth_path,
                                                              app.config['SERVER_HOSTNAME'])
  headers['WWW-Authenticate'] = authenticate
  headers['Docker-Distribution-API-Version'] = 'registry/2.0'
  return headers


def identity_from_bearer_token(bearer_token, max_signed_s, public_key):
  """ Process a bearer token and return the loaded identity, or raise InvalidJWTException if an
      identity could not be loaded. Expects tokens and grants in the format of the Docker registry
      v2 auth spec: https://docs.docker.com/registry/spec/auth/token/
  """
  logger.debug('Validating auth header: %s', bearer_token)

  # Extract the jwt token from the header
  match = TOKEN_REGEX.match(bearer_token)
  if match is None:
    raise InvalidJWTException('Invalid bearer token format')

  encoded = match.group(1)
  logger.debug('encoded JWT: %s', encoded)

  # Load the JWT returned.
  try:
    expected_issuer = app.config['JWT_AUTH_TOKEN_ISSUER']
    audience = app.config['SERVER_HOSTNAME']
    max_exp = strictjwt.exp_max_s_option(max_signed_s)
    payload = strictjwt.decode(encoded, public_key, algorithms=['RS256'], audience=audience,
                               issuer=expected_issuer, options=max_exp)
  except strictjwt.InvalidTokenError:
    logger.exception('Invalid token reason')
    raise InvalidJWTException('Invalid token')

  if not 'sub' in payload:
    raise InvalidJWTException('Missing sub field in JWT')

  loaded_identity = Identity(payload['sub'], 'signed_jwt')

  # Process the grants from the payload
  if 'access' in payload:
    try:
      validate(payload['access'], ACCESS_SCHEMA)
    except ValidationError:
      logger.exception('We should not be minting invalid credentials')
      raise InvalidJWTException('Token contained invalid or malformed access grants')

    for grant in payload['access']:
      namespace, repo_name = parse_namespace_repository(grant['name'])

      if 'push' in grant['actions']:
        loaded_identity.provides.add(repository_write_grant(namespace, repo_name))
      elif 'pull' in grant['actions']:
        loaded_identity.provides.add(repository_read_grant(namespace, repo_name))

  default_context = {
    'kind': 'anonymous'
  }

  if payload['sub'] != ANONYMOUS_SUB:
    default_context = {
      'kind': 'user',
      'user': payload['sub'],
    }

  return loaded_identity, payload.get('context', default_context)


@lru_cache(maxsize=1)
def load_public_key(certificate_file_path):
  with open(certificate_file_path) as cert_file:
    cert_obj = load_pem_x509_certificate(cert_file.read(), default_backend())
    return cert_obj.public_key()


def process_registry_jwt_auth(func):
  @wraps(func)
  def wrapper(*args, **kwargs):
    logger.debug('Called with params: %s, %s', args, kwargs)
    auth = request.headers.get('authorization', '').strip()
    if auth:
      max_signature_seconds = app.config.get('JWT_AUTH_MAX_FRESH_S', 3660)
      certificate_file_path = app.config['JWT_AUTH_CERTIFICATE_PATH']
      public_key = load_public_key(certificate_file_path)

      try:
        extracted_identity, context = identity_from_bearer_token(auth, max_signature_seconds,
                                                                 public_key)

        identity_changed.send(app, identity=extracted_identity)
        set_grant_context(context)
        logger.debug('Identity changed to %s', extracted_identity.id)
      except InvalidJWTException as ije:
        abort(401, message=ije.message, headers=get_auth_headers())
    else:
      logger.debug('No auth header.')

    return func(*args, **kwargs)
  return wrapper