woodruffw / pyrage

Python bindings for rage (age in Rust)
https://pypi.org/project/pyrage/
MIT License
52 stars 7 forks source link

Add ASCII Armored Option [Feature Request] #45

Open furechan opened 11 months ago

furechan commented 11 months ago

Hi, I couldn't find a way to generate an ascii armored output with the encrypt function. Is this something that could be added easily ?

Thanks for the great package!

woodruffw commented 11 months ago

Is this something that could be added easily ?

I think so -- rage supports armoring and dearmoring, so we'll probably just want an armored: bool = False kwarg on both encryption and decryption.

Patches are welcome for this :slightly_smiling_face:

jirib commented 10 months ago

I'm playing a bit with pysimplegui to do a GUI based on pyrage, and armor feature would be nice to have a possibility to encrypt a multiline form.

vikanezrimaya commented 7 months ago

I may or may not try my hand at this after finishing up #56.

Alchemyst0x commented 2 days ago

For fun, I implemented this in Python. It works, and I tried my best to ensure that it is strictly RFC-compliant and everything, but I am no expert so feel free to point out any problems in my implementation.

This is the relevant portion of the module I wrote:

@dataclass
class Armored:
    """RFC-compliant ASCII Armor implementation for age encryption."""

    PEM_HEADER = '-----BEGIN AGE ENCRYPTED FILE-----'
    PEM_FOOTER = '-----END AGE ENCRYPTED FILE-----'

    PEM_RE = re.compile(
        rf'^{PEM_HEADER}\n' r'([A-Za-z0-9+/=\n]+)' rf'\n{PEM_FOOTER}$',
    )
    B64_LINE_RE = re.compile(
        r'(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=|[A-Za-z0-9+/]{4})?'
    )

    @property
    def armored_data(self) -> str:
        return self._armored_data

    @property
    def dearmored_data(self) -> bytes:
        return self._dearmored_data

    def __init__(self, data: bytes | str) -> None:
        if isinstance(data, bytes):
            self._armored_data = self._armor(data)
            self._dearmored_data = self._dearmor(self._armored_data)
        elif isinstance(data, str):
            self._dearmored_data = self._dearmor(data)
            self._armored_data = self._armor(self._dearmored_data)
        else:
            raise TypeError

    def _decode_b64_strict(self, b64_data: str) -> bytes:
        while '\r\n' in b64_data:
            b64_data = b64_data.replace('\r\n', '\n')
        while '\r' in b64_data:
            b64_data = b64_data.replace('\r', '\n')

        b64_lines = b64_data.split('\n')
        for idx, line in enumerate(b64_lines):
            if idx < len(b64_lines) - 1:
                if len(line) != 64:
                    raise ValueError(f'Line {idx+1} length is not 64 characters.')
            elif len(line) > 64:
                raise ValueError('Final line length exceeds 64 characters.')

        b64_str = ''.join(b64_lines)
        if not re.fullmatch(self.B64_LINE_RE, b64_str):
            raise ValueError('Invalid Base64 encoding detected.')

        try:
            decoded_data = binascii.a2b_base64(b64_str, strict_mode=True)
        except binascii.Error as exc:
            raise ValueError('Base64 decoding error: ' + str(exc)) from exc
        return decoded_data

    def _armor(self, data: bytes) -> str:
        b64_encoded = binascii.b2a_base64(data, newline=False).decode('ascii')
        b64_lines = [b64_encoded[i : i + 64] for i in range(0, len(b64_encoded), 64)]
        return '\n'.join([self.PEM_HEADER, *b64_lines, self.PEM_FOOTER])

    def _dearmor(self, pem_data: str) -> bytes:
        pem_data = pem_data.strip()
        match = re.fullmatch(self.PEM_RE, pem_data)
        if not match:
            raise ValueError('Invalid PEM format or extra data found.')
        b64_data = match.group(1)
        return self._decode_b64_strict(b64_data)

Figured I'd share this here in case it was helpful to anyone else.