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
|
from contextlib import contextmanager
try:
from unittest.mock import patch
except ImportError:
from mock import patch
class FileReadWrite(object):
"""Mock an opened file that can be read and written to."""
def __init__(self):
self._content = b''
self._mode = 'r'
def read(self):
if self._mode == 'r':
if not isinstance(self._content, str):
return self._content.decode()
return self._content
if isinstance(self._content, str):
return self._content.encode('utf-8')
return self._content
def write(self, content):
self._content = content
PRIVATE_KEY = FileReadWrite()
PUBLIC_KEY = FileReadWrite()
@contextmanager
def open_priv_pub(infile, mode='r'):
try:
if infile.endswith('.pub'):
PUBLIC_KEY._mode = mode
yield PUBLIC_KEY
else:
PRIVATE_KEY._mode = mode
yield PRIVATE_KEY
finally:
pass
|
