#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from ... import Unit, RefineryPartialResult
from ....lib.argformats import multibin

from Crypto.Cipher import PKCS1_OAEP, PKCS1_v1_5
from Crypto.PublicKey import RSA


class rsa(Unit):
    """
    Implements single block RSA encryption and decryption. This unit can be used to encrypt
    and decrypt blocks generated by openssl's `rsautl` tool when using the mode `-verify`.
    When it is executed with a public key for decryption or with a private key for encryption,
    it will perform a raw RSA operation. The result of these operations are (un)padded using
    EMSA-PKCS1-v1_5.
    """
    def interface(self, argp):
        argp.add_argument('key', type=multibin, help='RSA key in PEM or DER format.')
        return super().interface(argp)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.key = RSA.import_key(self.args.key)

    @property
    def blocksize(self) -> int:
        return self.key.size_in_bytes()

    def process(self, data):
        if not self.key.has_private():
            if len(data) > self.blocksize:
                raise ValueError(F'This key can only handle messages of size {self.blocksize}.')
            self.log_debug('attempting to inverse decrypt as raw blob')
            msg = int.from_bytes(data, byteorder='big')
            dec = pow(msg, self.key.e, self.key.n).to_bytes(self.blocksize, byteorder='big')
            if dec.startswith(B'\x00\x01'):
                pos = dec.find(B'\0', 2)
                if pos > 0:
                    pad = dec[2:pos]
                    if all(b == 0xFF for b in pad):
                        return dec[pos + 1:]
            raise RefineryPartialResult(
                'The decrypted message did not have a correct EMSA-PKCS1-v1_5 padding.', dec)
        result = PKCS1_v1_5.new(self.key).decrypt(data, None)
        if result is not None:
            return result
        return PKCS1_OAEP.new(self.key).decrypt(data)

    def reverse(self, data):
        if self.key.has_private():
            if len(data) > self.blocksize - 11:
                raise ValueError(F'This key can only encrypt messages of size at most {self.blocksize - 11}.')
            self.log_debug('attempting to inverse encrypt as raw blob')
            pad = self.blocksize - len(data) - 3
            msg = B'\x00\x01' + B'\xFF' * pad + B'\x00' + data
            msg = int.from_bytes(msg, byteorder='big')
            enc = pow(msg, self.key.d, self.key.n).to_bytes(self.blocksize, byteorder='big')
            return enc
        return PKCS1_OAEP.new(self.key).encrypt(data)
