-
Notifications
You must be signed in to change notification settings - Fork 59
/
keygen_stub.py
40 lines (30 loc) · 949 Bytes
/
keygen_stub.py
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
from contextlib import contextmanager
from unittest.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