# passerelle - uniform access to multiple data sources and services
# Copyright (C) 2020  Entr'ouvert
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import base64
import binascii
import json
import os
from urllib.parse import urljoin
from uuid import uuid4

from Cryptodome.Cipher import AES, PKCS1_OAEP
from Cryptodome.Hash import SHA512
from Cryptodome.PublicKey import RSA
from Cryptodome.Random import get_random_bytes
from django.core.exceptions import ValidationError
from django.core.files.storage import default_storage
from django.db import models
from django.http import HttpResponse
from django.utils.translation import gettext_lazy as _

from passerelle.base.models import BaseResource
from passerelle.utils.api import endpoint
from passerelle.utils.files import atomic_write
from passerelle.utils.jsonresponse import APIError

FILE_SCHEMA = {
    "$schema": "http://json-schema.org/draft-04/schema#",
    "title": "File to encrypt",
    "description": "",
    "type": "object",
    "required": ["file"],
    "properties": {
        "file": {
            "type": "object",
            "required": ["filename", "content_type", "content"],
            "properties": {
                "filename": {"type": "string"},
                "content_type": {"type": "string"},
                "content": {"type": "string"},
            },
        }
    },
}


# encrypt and decrypt are borrowed from
# https://www.pycryptodome.org/en/latest/src/examples.html#encrypt-data-with-rsa


def write_encrypt(out_file, data, key_pem):
    public_key = RSA.import_key(key_pem)
    session_key = get_random_bytes(16)
    # Encrypt the session key with the public RSA key
    cipher_rsa = PKCS1_OAEP.new(public_key, hashAlgo=SHA512)
    enc_session_key = cipher_rsa.encrypt(session_key)
    # Encrypt the data with the AES session key
    cipher_aes = AES.new(session_key, AES.MODE_EAX)
    ciphertext, tag = cipher_aes.encrypt_and_digest(data)
    # Store in out_file
    out_file.write(enc_session_key)
    out_file.write(cipher_aes.nonce)
    out_file.write(tag)
    out_file.write(ciphertext)


def read_decrypt(in_file, key_pem):
    private_key = RSA.import_key(key_pem)
    # Get crypt elements from in_file
    enc_session_key = in_file.read(private_key.size_in_bytes())
    nonce = in_file.read(16)
    tag = in_file.read(16)
    ciphertext = in_file.read()
    # Decrypt the session key with the private RSA key
    cipher_rsa = PKCS1_OAEP.new(private_key, hashAlgo=SHA512)
    session_key = cipher_rsa.decrypt(enc_session_key)
    # Decrypt the data with the AES session key
    cipher_aes = AES.new(session_key, AES.MODE_EAX, nonce)
    decrypted = cipher_aes.decrypt_and_verify(ciphertext, tag)
    return decrypted


def makedir(dir_name):
    if not os.path.exists(dir_name):
        if default_storage.directory_permissions_mode:
            d_umask = os.umask(0)
            try:
                os.makedirs(dir_name, mode=default_storage.directory_permissions_mode)
            except OSError:
                pass
            finally:
                os.umask(d_umask)
        else:
            os.makedirs(dir_name)


def validate_rsa_key(key):
    try:
        RSA.import_key(key)
    except ValueError as ex:
        raise ValidationError(_('invalid RSA key (%s)') % ex)


class Cryptor(BaseResource):
    public_key = models.TextField(
        blank=True, verbose_name=_('Encryption RSA public key (PEM format)'), validators=[validate_rsa_key]
    )
    private_key = models.TextField(
        blank=True, verbose_name=_('Decryption RSA private key (PEM format)'), validators=[validate_rsa_key]
    )
    redirect_url_base = models.URLField(
        max_length=256,
        blank=True,
        verbose_name=_('Base URL of decryption system'),
        help_text=_('Base URL for redirect, empty for local'),
    )

    category = _('Misc')

    class Meta:
        verbose_name = _('Encryption / Decryption')

    def get_redirect_url_base_display(self):
        if self.redirect_url_base:
            return _('defined')  # don't show it, can be sensitive
        return _('this file-decrypt endpoint')

    def get_filename(self, uuid, create=False):
        dirname = os.path.join(
            default_storage.path(self.get_connector_slug()), self.slug, uuid[0:2], uuid[2:4]
        )
        if create:
            makedir(dirname)
        filename = os.path.join(dirname, uuid)
        return filename

    @endpoint(
        name='file-encrypt',
        perm='can_encrypt',
        description=_('Encrypt a file'),
        post={
            'description': _('File to encrypt'),
            'request_body': {'schema': {'application/json': FILE_SCHEMA}},
        },
    )
    def file_encrypt(self, request, post_data):
        if not self.public_key:
            raise APIError('missing public key')
        try:
            data = base64.b64decode(post_data['file']['content'])
        except (TypeError, binascii.Error):
            raise APIError('invalid base64 string', http_status=400)

        filename = post_data['file']['filename']
        content_type = post_data['file']['content_type']
        cfile = CryptedFile(resource=self, filename=filename, content_type=content_type)
        cfile.save()

        uuid = str(cfile.uuid)  # get string representation of UUID object

        if self.redirect_url_base:
            redirect_url_base = self.redirect_url_base
        else:
            redirect_url_base = request.build_absolute_uri('%sfile-decrypt/' % (self.get_absolute_url(),))
        redirect_url = urljoin(redirect_url_base, uuid)

        content_filename = self.get_filename(uuid, create=True)
        metadata_filename = '%s.json' % content_filename
        metadata = {
            'filename': cfile.filename,
            'content_type': cfile.content_type,
            'creation_timestamp': cfile.creation_timestamp.isoformat(),
            'redirect_url': redirect_url,
        }

        tmp_dir = os.path.join(default_storage.path(self.get_connector_slug()), self.slug, 'tmp')
        with atomic_write(content_filename, dir=tmp_dir) as fd:
            write_encrypt(fd, data, self.public_key)
        with atomic_write(metadata_filename, dir=tmp_dir, mode='w') as fd:
            json.dump(metadata, fd, indent=2)

        return {'data': metadata}

    @endpoint(
        name='file-decrypt',
        perm='can_decrypt',
        description=_('Decrypt a file'),
        pattern=r'(?P<uuid>[\w-]+)$',
        example_pattern='{uuid}/',
        parameters={
            'uuid': {
                'description': _('File identifier'),
                'example_value': '12345678-abcd-4321-abcd-123456789012',
            },
        },
    )
    def file_decrypt(self, request, uuid):
        if not self.private_key:
            raise APIError('missing private key')
        content_filename = self.get_filename(uuid, create=False)
        metadata_filename = '%s.json' % content_filename
        if not os.path.exists(metadata_filename):
            raise APIError('unknown uuid', http_status=404)

        with open(content_filename, 'rb') as fd:
            content = read_decrypt(fd, self.private_key)

        with open(metadata_filename) as fd:
            metadata = json.load(fd)
        filename = metadata.get('filename')
        content_type = metadata.get('content_type')

        response = HttpResponse(content_type=content_type)
        response['Content-Disposition'] = 'inline; filename="%s"' % filename
        response.write(content)
        return response


class CryptedFile(models.Model):
    resource = models.ForeignKey(Cryptor, on_delete=models.PROTECT)
    uuid = models.UUIDField(primary_key=True, default=uuid4, editable=False)
    filename = models.CharField(max_length=512, blank=False)
    content_type = models.CharField(max_length=128)
    creation_timestamp = models.DateTimeField(auto_now_add=True)
