1# coding: utf-8
2from __future__ import unicode_literals, division, absolute_import, print_function
3
4import unittest
5import sys
6import os
7
8from asn1crypto import pem, util
9
10from .unittest_data import data_decorator, data
11from ._unittest_compat import patch
12
13patch()
14
15if sys.version_info < (3,):
16    byte_cls = str
17    num_cls = long  # noqa
18else:
19    byte_cls = bytes
20    num_cls = int
21
22
23tests_root = os.path.dirname(__file__)
24fixtures_dir = os.path.join(tests_root, 'fixtures')
25
26
27@data_decorator
28class PEMTests(unittest.TestCase):
29
30    @staticmethod
31    def detect_files():
32        return (
33            (
34                'keys/test-der.crt',
35                False
36            ),
37            (
38                'keys/test-inter-der.crt',
39                False
40            ),
41            (
42                'keys/test-third-der.crt',
43                False
44            ),
45            (
46                'keys/test.crt',
47                True
48            ),
49            (
50                'keys/test-inter.crt',
51                True
52            ),
53            (
54                'keys/test-third.crt',
55                True
56            ),
57        )
58
59    @data('detect_files')
60    def detect(self, relative_path, is_pem):
61        with open(os.path.join(fixtures_dir, relative_path), 'rb') as f:
62            byte_string = f.read()
63        self.assertEqual(is_pem, pem.detect(byte_string))
64
65    @staticmethod
66    def unarmor_armor_files():
67        return (
68            (
69                'keys/test.crt',
70                'keys/test-der.crt',
71                'CERTIFICATE',
72                {}
73            ),
74            (
75                'keys/test-inter.crt',
76                'keys/test-inter-der.crt',
77                'CERTIFICATE',
78                {}
79            ),
80            (
81                'keys/test-third.crt',
82                'keys/test-third-der.crt',
83                'CERTIFICATE',
84                {}
85            ),
86            (
87                'keys/test-pkcs8.key',
88                'keys/test-pkcs8-der.key',
89                'PRIVATE KEY',
90                {}
91            ),
92            (
93                'test-third.csr',
94                'test-third-der.csr',
95                'CERTIFICATE REQUEST',
96                {}
97            ),
98            (
99                'keys/test-aes128.key',
100                'keys/test-aes128-der.key',
101                'RSA PRIVATE KEY',
102                util.OrderedDict([
103                    ('Proc-Type', '4,ENCRYPTED'),
104                    ('DEK-Info', 'AES-128-CBC,01F6EE04516C912788B11BD7377626C2')
105                ])
106            ),
107        )
108
109    @data('unarmor_armor_files')
110    def unarmor(self, relative_path, expected_bytes_filename, expected_type_name, expected_headers):
111        with open(os.path.join(fixtures_dir, relative_path), 'rb') as f:
112            byte_string = f.read()
113
114        type_name, headers, decoded_bytes = pem.unarmor(byte_string)
115        self.assertEqual(expected_type_name, type_name)
116        self.assertEqual(expected_headers, headers)
117        with open(os.path.join(fixtures_dir, expected_bytes_filename), 'rb') as f:
118            expected_bytes = f.read()
119            self.assertEqual(expected_bytes, decoded_bytes)
120
121    def test_unarmor_multiple(self):
122        data = self.unarmor_armor_files()
123        input_data = b''
124        der_data = []
125        for pem_file, der_file in ((data[0][0], data[0][1]), (data[1][0], data[1][1])):
126            with open(os.path.join(fixtures_dir, pem_file), 'rb') as f:
127                input_data += f.read() + b'\n'
128            with open(os.path.join(fixtures_dir, der_file), 'rb') as f:
129                der_data.append(f.read())
130        i = 0
131        for name, headers, der_bytes in pem.unarmor(input_data, True):
132            self.assertEqual('CERTIFICATE', name)
133            self.assertEqual({}, headers)
134            self.assertEqual(der_data[i], der_bytes)
135            i += 1
136        self.assertEqual(2, i)
137
138    @data('unarmor_armor_files')
139    def armor(self, expected_bytes_filename, relative_path, type_name, headers):
140        with open(os.path.join(fixtures_dir, relative_path), 'rb') as f:
141            byte_string = f.read()
142
143        encoded_bytes = pem.armor(type_name, byte_string, headers=headers)
144        with open(os.path.join(fixtures_dir, expected_bytes_filename), 'rb') as f:
145            expected_bytes = f.read()
146            # In case a user on Windows has CRLF translation on in Git.
147            # Ran into this with the GitHub Actions Windows environments.
148            expected_bytes = expected_bytes.replace(b'\r\n', b'\n')
149            self.assertEqual(expected_bytes, encoded_bytes)
150
151    def test_armor_wrong_type(self):
152        with self.assertRaisesRegex(TypeError, 'type_name must be a unicode string'):
153            pem.armor(b'CERTIFICATE', b'')
154
155    def test_armor_wrong_type2(self):
156        with self.assertRaisesRegex(TypeError, 'der_bytes must be a byte string'):
157            pem.armor('CERTIFICATE', '')
158
159    def test_detect_wrong_type(self):
160        with self.assertRaisesRegex(TypeError, 'byte_string must be a byte string'):
161            pem.detect('CERTIFICATE')
162