# openssl genrsa -out private.pem 2048
# openssl rsa -in private.pem -pubout -out public.pem
# pip install cryptography
import sys
from argparse import ArgumentParser
from base64 import b64encode
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import (
    hashes,
    serialization,
    )
from cryptography.hazmat.primitives.asymmetric import padding


backend = default_backend()
hash_sha256 = hashes.SHA256()


def sign(private_key: bytes, data: bytes) -> bytes:
    key = serialization.load_pem_private_key(
            private_key, password=None, backend=backend)
    mgf = padding.MGF1(hash_sha256)
    padding_obj = padding.PSS(mgf=mgf, salt_length=padding.PSS.MAX_LENGTH)
    return key.sign(data, padding_obj, hash_sha256)


def verify(public_key: bytes, data_signed: bytes, data: bytes):
    key = serialization.load_pem_public_key(public_key, backend=backend)
    mgf = padding.MGF1(hash_sha256)
    padding_obj = padding.PSS(mgf=mgf, salt_length=padding.PSS.MAX_LENGTH)
    key.verify(data_signed, data, padding_obj, hash_sha256)


def sign_without_salt(private_key: bytes, data: bytes) -> bytes:
    key = serialization.load_pem_private_key(
            private_key, password=None, backend=backend)
    return key.sign(data, padding.PKCS1v15(), hash_sha256)


def verify_without_salt(public_key: bytes, data_signed: bytes, data: bytes):
    key = serialization.load_pem_public_key(public_key, backend=backend)
    key.verify(data_signed, data, padding.PKCS1v15(), hash_sha256)


def main(argv=sys.argv[1:]):
    pars = ArgumentParser()
    pars.add_argument('--private-file', required=True)
    pars.add_argument('--data-file', required=True)
    pars.add_argument('--without-salt', action='store_true')
    pars.add_argument('--public-file')
    option = pars.parse_args(argv)

    with open(option.private_file, 'rb') as f:
        private_key = f.read()

    with open(option.data_file, 'rb') as f:
        data = f.read()

    sign_func = option.without_salt and sign_without_salt or sign
    data_signed = sign_func(private_key, data)
    data_b64 = b64encode(data_signed).decode('utf-8')
    print(data_b64)
    if not option.public_file:
        return
    with open(option.public_file, 'rb') as f:
        public_key = f.read()
    verify_func = option.without_salt and verify_without_salt or verify
    verify_func(public_key, data_signed, data)


if __name__ == '__main__':
    main()