summaryrefslogtreecommitdiff
path: root/rsa/pem.py
blob: 0875ec32c4d0eac4e3be2037a6060ffb9556bd8b (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
# -*- coding: utf-8 -*-
#
#  Copyright 2011 Sybren A. Stüvel <sybren@stuvel.eu>
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

'''Functions that load and write PEM-encoded files.'''

import base64
from rsa._compat import b, is_bytes

def _markers(pem_marker):
    '''Returns the start and end PEM markers

    >>> _markers('RSA PRIVATE KEY')
    ('-----BEGIN RSA PRIVATE KEY-----', '-----END RSA PRIVATE KEY-----')

    '''

    if is_bytes(pem_marker):
        pem_marker = pem_marker.decode('utf-8')

    return (b('-----BEGIN %s-----' % pem_marker),
            b('-----END %s-----' % pem_marker))

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.

    '''

    (pem_start, pem_end) = _markers(pem_marker)

    pem_lines = []
    in_pem_part = False

    for line in contents.splitlines():
        line = line.strip()

        # Skip empty lines
        if not line:
            continue

        # Handle start marker
        if line == pem_start:
            if in_pem_part:
                raise ValueError('Seen start marker "%s" twice' % pem_start)

            in_pem_part = True
            continue

        # Skip stuff before first marker
        if not in_pem_part:
            continue

        # Handle end marker
        if in_pem_part and line == pem_end:
            in_pem_part = False
            break

        # Load fields
        if b(':') in line:
            continue

        pem_lines.append(line)

    # Do some sanity checks
    if not pem_lines:
        raise ValueError('No PEM start marker "%s" found' % pem_start)

    if in_pem_part:
        raise ValueError('No PEM end marker "%s" found' % pem_end)

    # Base64-decode the contents
    pem = b('').join(pem_lines)
    return base64.decodestring(pem)


def save_pem(contents, pem_marker):
    '''Saves a PEM file.

    @param contents: the contents to encode in PEM format
    @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-encoded content between the start and end markers.

    '''

    (pem_start, pem_end) = _markers(pem_marker)

    b64 = base64.encodestring(contents).replace(b('\n'), b(''))
    pem_lines = [pem_start]
    
    for block_start in range(0, len(b64), 64):
        block = b64[block_start:block_start + 64]
        pem_lines.append(block)

    pem_lines.append(pem_end)
    pem_lines.append(b(''))

    return b('\n').join(pem_lines)