You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

132 lines
3.9 KiB

  1. # Copyright 2011 Sybren A. Stüvel <sybren@stuvel.eu>
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # https://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. """Functions that load and write PEM-encoded files."""
  15. import base64
  16. import typing
  17. # Should either be ASCII strings or bytes.
  18. FlexiText = typing.Union[str, bytes]
  19. def _markers(pem_marker: FlexiText) -> typing.Tuple[bytes, bytes]:
  20. """
  21. Returns the start and end PEM markers, as bytes.
  22. """
  23. if not isinstance(pem_marker, bytes):
  24. pem_marker = pem_marker.encode('ascii')
  25. return (b'-----BEGIN ' + pem_marker + b'-----',
  26. b'-----END ' + pem_marker + b'-----')
  27. def _pem_lines(contents: bytes, pem_start: bytes, pem_end: bytes) -> typing.Iterator[bytes]:
  28. """Generator over PEM lines between pem_start and pem_end."""
  29. in_pem_part = False
  30. seen_pem_start = False
  31. for line in contents.splitlines():
  32. line = line.strip()
  33. # Skip empty lines
  34. if not line:
  35. continue
  36. # Handle start marker
  37. if line == pem_start:
  38. if in_pem_part:
  39. raise ValueError('Seen start marker "%r" twice' % pem_start)
  40. in_pem_part = True
  41. seen_pem_start = True
  42. continue
  43. # Skip stuff before first marker
  44. if not in_pem_part:
  45. continue
  46. # Handle end marker
  47. if in_pem_part and line == pem_end:
  48. in_pem_part = False
  49. break
  50. # Load fields
  51. if b':' in line:
  52. continue
  53. yield line
  54. # Do some sanity checks
  55. if not seen_pem_start:
  56. raise ValueError('No PEM start marker "%r" found' % pem_start)
  57. if in_pem_part:
  58. raise ValueError('No PEM end marker "%r" found' % pem_end)
  59. def load_pem(contents: FlexiText, pem_marker: FlexiText) -> bytes:
  60. """Loads a PEM file.
  61. :param contents: the contents of the file to interpret
  62. :param pem_marker: the marker of the PEM content, such as 'RSA PRIVATE KEY'
  63. when your file has '-----BEGIN RSA PRIVATE KEY-----' and
  64. '-----END RSA PRIVATE KEY-----' markers.
  65. :return: the base64-decoded content between the start and end markers.
  66. @raise ValueError: when the content is invalid, for example when the start
  67. marker cannot be found.
  68. """
  69. # We want bytes, not text. If it's text, it can be converted to ASCII bytes.
  70. if not isinstance(contents, bytes):
  71. contents = contents.encode('ascii')
  72. (pem_start, pem_end) = _markers(pem_marker)
  73. pem_lines = [line for line in _pem_lines(contents, pem_start, pem_end)]
  74. # Base64-decode the contents
  75. pem = b''.join(pem_lines)
  76. return base64.standard_b64decode(pem)
  77. def save_pem(contents: bytes, pem_marker: FlexiText) -> bytes:
  78. """Saves a PEM file.
  79. :param contents: the contents to encode in PEM format
  80. :param pem_marker: the marker of the PEM content, such as 'RSA PRIVATE KEY'
  81. when your file has '-----BEGIN RSA PRIVATE KEY-----' and
  82. '-----END RSA PRIVATE KEY-----' markers.
  83. :return: the base64-encoded content between the start and end markers, as bytes.
  84. """
  85. (pem_start, pem_end) = _markers(pem_marker)
  86. b64 = base64.standard_b64encode(contents).replace(b'\n', b'')
  87. pem_lines = [pem_start]
  88. for block_start in range(0, len(b64), 64):
  89. block = b64[block_start:block_start + 64]
  90. pem_lines.append(block)
  91. pem_lines.append(pem_end)
  92. pem_lines.append(b'')
  93. return b'\n'.join(pem_lines)