diff options
Diffstat (limited to 'rsa/pem.py')
| -rw-r--r-- | rsa/pem.py | 70 |
1 files changed, 38 insertions, 32 deletions
@@ -1,5 +1,3 @@ -# -*- coding: utf-8 -*- -# # Copyright 2011 Sybren A. Stüvel <sybren@stuvel.eu> # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -17,45 +15,29 @@ """Functions that load and write PEM-encoded files.""" import base64 +import typing -from rsa._compat import is_bytes, range +# Should either be ASCII strings or bytes. +FlexiText = typing.Union[str, bytes] -def _markers(pem_marker): +def _markers(pem_marker: FlexiText) -> typing.Tuple[bytes, bytes]: """ Returns the start and end PEM markers, as bytes. """ - if not is_bytes(pem_marker): + if not isinstance(pem_marker, bytes): pem_marker = pem_marker.encode('ascii') return (b'-----BEGIN ' + pem_marker + b'-----', b'-----END ' + pem_marker + b'-----') -def load_pem(contents, pem_marker): - """Loads a PEM file. - - :param contents: the contents of the file to interpret - :param pem_marker: the marker of the PEM content, such as 'RSA PRIVATE KEY' - when your file has '-----BEGIN RSA PRIVATE KEY-----' and - '-----END RSA PRIVATE KEY-----' markers. - - :return: the base64-decoded content between the start and end markers. - - @raise ValueError: when the content is invalid, for example when the start - marker cannot be found. - - """ - - # We want bytes, not text. If it's text, it can be converted to ASCII bytes. - if not is_bytes(contents): - contents = contents.encode('ascii') - - (pem_start, pem_end) = _markers(pem_marker) +def _pem_lines(contents: bytes, pem_start: bytes, pem_end: bytes) -> typing.Iterator[bytes]: + """Generator over PEM lines between pem_start and pem_end.""" - pem_lines = [] in_pem_part = False + seen_pem_start = False for line in contents.splitlines(): line = line.strip() @@ -67,9 +49,10 @@ def load_pem(contents, pem_marker): # Handle start marker if line == pem_start: if in_pem_part: - raise ValueError('Seen start marker "%s" twice' % pem_start) + raise ValueError('Seen start marker "%r" twice' % pem_start) in_pem_part = True + seen_pem_start = True continue # Skip stuff before first marker @@ -85,21 +68,44 @@ def load_pem(contents, pem_marker): if b':' in line: continue - pem_lines.append(line) + yield line # Do some sanity checks - if not pem_lines: - raise ValueError('No PEM start marker "%s" found' % pem_start) + if not seen_pem_start: + raise ValueError('No PEM start marker "%r" found' % pem_start) if in_pem_part: - raise ValueError('No PEM end marker "%s" found' % pem_end) + raise ValueError('No PEM end marker "%r" found' % pem_end) + + +def load_pem(contents: FlexiText, pem_marker: FlexiText) -> bytes: + """Loads a PEM file. + + :param contents: the contents of the file to interpret + :param pem_marker: the marker of the PEM content, such as 'RSA PRIVATE KEY' + when your file has '-----BEGIN RSA PRIVATE KEY-----' and + '-----END RSA PRIVATE KEY-----' markers. + + :return: the base64-decoded content between the start and end markers. + + @raise ValueError: when the content is invalid, for example when the start + marker cannot be found. + + """ + + # We want bytes, not text. If it's text, it can be converted to ASCII bytes. + if not isinstance(contents, bytes): + contents = contents.encode('ascii') + + (pem_start, pem_end) = _markers(pem_marker) + pem_lines = [line for line in _pem_lines(contents, pem_start, pem_end)] # Base64-decode the contents pem = b''.join(pem_lines) return base64.standard_b64decode(pem) -def save_pem(contents, pem_marker): +def save_pem(contents: bytes, pem_marker: FlexiText) -> bytes: """Saves a PEM file. :param contents: the contents to encode in PEM format |