import base64
import json
from typing import Union
from Crypto.Signature import PKCS1_v1_5
from Crypto.PublicKey import RSA
from Crypto.Hash import SHA256


def create_signature(private_key: str, data: dict) -> str:
    """
    Create sign with RSA private key
    :param private_key: str
    :param data: dict
    :return: str
    """
    message = json.dumps(data)
    message = message.encode('utf8')

    signer = PKCS1_v1_5.new(RSA.importKey(private_key))
    digest = SHA256.new()
    digest.update(message)
    result = base64.b64encode(signer.sign(digest))
    return result.decode()


def validate_signature(public_key: str, signature: str, message: Union[str, bytes]) -> bool:
    """
    Check if signature is valid with RSA public key
    :param public_key: str
    :param signature: str
    :param message: str or bytes
    :return: bool
    """
    if isinstance(message, str):
        message = message.encode('utf8')

    signer = PKCS1_v1_5.new(RSA.importKey(public_key))
    digest = SHA256.new()
    digest.update(message)

    try:
        if signer.verify(digest, base64.b64decode(signature)):
            return True
    except Exception:
        return False

    return False