1# Copyright 2015 Google Inc. All rights reserved.
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#      http://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"""Helper functions for commonly used utilities."""
15
16import base64
17import json
18
19import six
20
21
22def _parse_pem_key(raw_key_input):
23    """Identify and extract PEM keys.
24
25    Determines whether the given key is in the format of PEM key, and extracts
26    the relevant part of the key if it is.
27
28    Args:
29        raw_key_input: The contents of a private key file (either PEM or
30                       PKCS12).
31
32    Returns:
33        string, The actual key if the contents are from a PEM file, or
34        else None.
35    """
36    offset = raw_key_input.find(b'-----BEGIN ')
37    if offset != -1:
38        return raw_key_input[offset:]
39
40
41def _json_encode(data):
42    return json.dumps(data, separators=(',', ':'))
43
44
45def _to_bytes(value, encoding='ascii'):
46    """Converts a string value to bytes, if necessary.
47
48    Unfortunately, ``six.b`` is insufficient for this task since in
49    Python2 it does not modify ``unicode`` objects.
50
51    Args:
52        value: The string/bytes value to be converted.
53        encoding: The encoding to use to convert unicode to bytes. Defaults
54                  to "ascii", which will not allow any characters from ordinals
55                  larger than 127. Other useful values are "latin-1", which
56                  which will only allows byte ordinals (up to 255) and "utf-8",
57                  which will encode any unicode that needs to be.
58
59    Returns:
60        The original value converted to bytes (if unicode) or as passed in
61        if it started out as bytes.
62
63    Raises:
64        ValueError if the value could not be converted to bytes.
65    """
66    result = (value.encode(encoding)
67              if isinstance(value, six.text_type) else value)
68    if isinstance(result, six.binary_type):
69        return result
70    else:
71        raise ValueError('{0!r} could not be converted to bytes'.format(value))
72
73
74def _from_bytes(value):
75    """Converts bytes to a string value, if necessary.
76
77    Args:
78        value: The string/bytes value to be converted.
79
80    Returns:
81        The original value converted to unicode (if bytes) or as passed in
82        if it started out as unicode.
83
84    Raises:
85        ValueError if the value could not be converted to unicode.
86    """
87    result = (value.decode('utf-8')
88              if isinstance(value, six.binary_type) else value)
89    if isinstance(result, six.text_type):
90        return result
91    else:
92        raise ValueError(
93            '{0!r} could not be converted to unicode'.format(value))
94
95
96def _urlsafe_b64encode(raw_bytes):
97    raw_bytes = _to_bytes(raw_bytes, encoding='utf-8')
98    return base64.urlsafe_b64encode(raw_bytes).rstrip(b'=')
99
100
101def _urlsafe_b64decode(b64string):
102    # Guard against unicode strings, which base64 can't handle.
103    b64string = _to_bytes(b64string)
104    padded = b64string + b'=' * (4 - len(b64string) % 4)
105    return base64.urlsafe_b64decode(padded)
106