#
# Copyright (c) 2021 the Hugging Face team.
#
# 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.#
import base64
import datetime

from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding

from app.core.config import EXPIRATION_MINUTES, SECRET_KEY
from app.models.experiment_join import HivemindAccess


PADDING = padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH)
HASH_ALGORITHM = hashes.SHA256()


def save_private_key(private_key):
    pem = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.TraditionalOpenSSL,
        encryption_algorithm=serialization.BestAvailableEncryption(f"{SECRET_KEY}".encode()),
    )
    return pem


def save_public_key(public_key):
    pem = public_key.public_bytes(
        encoding=serialization.Encoding.OpenSSH,
        format=serialization.PublicFormat.OpenSSH,
    )
    return pem


def load_private_key(string_in_db):
    private_key = serialization.load_pem_private_key(string_in_db, f"{SECRET_KEY}".encode())
    return private_key


def load_public_key(string_in_db):
    public_key = serialization.load_ssh_public_key(string_in_db)
    return public_key


def create_hivemind_access(peer_public_key: bytes, auth_server_private_key: bytes, username: str):
    current_time = datetime.datetime.utcnow()
    expiration_time = current_time + datetime.timedelta(minutes=EXPIRATION_MINUTES)

    private_key = load_private_key(auth_server_private_key)
    signature = private_key.sign(
        f"{username} {peer_public_key} {expiration_time}".encode(),
        PADDING,
        HASH_ALGORITHM,
    )
    signature = base64.b64encode(signature)

    hivemind_access = HivemindAccess(
        username=username,
        peer_public_key=peer_public_key,
        expiration_time=expiration_time,
        signature=signature,
    )
    return hivemind_access
