-
Notifications
You must be signed in to change notification settings - Fork 0
/
ssh_mac.py
132 lines (98 loc) · 2.91 KB
/
ssh_mac.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
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
124
125
126
127
128
129
130
131
132
"""
ssh mac 支持
"""
import abc
import hashlib
import hmac
import secrets
import typing as t
try:
import umac
umac_support = True
except ImportError:
umac_support = False
class MacInterface(abc.ABC):
length = 0
key_size = 0
seq_bytes = 4
is_etm = False
def __init__(self, key: bytes):
self.key = key
@abc.abstractmethod
def verify(self, message: bytes, mac: bytes) -> bool:
raise NotImplementedError("verify")
@abc.abstractmethod
def calculate(self, message: bytes) -> bytes:
raise NotImplementedError("calculate")
if umac_support:
class UmacImpl(MacInterface):
length = 8
key_size = 16
seq_bytes = 8
umac_cls = umac.Umac
def __init__(self, key: bytes):
super().__init__(key)
self.umac = self.umac_cls(self.key)
def verify(self, message: bytes, mac: bytes) -> bool:
nonce = message[:8]
message = message[8:]
got = self.umac.mac(message, nonce)
return got == mac
def calculate(self, message: bytes) -> bytes:
nonce = message[:8]
message = message[8:]
got = self.umac.mac(message, nonce)
return got
class UmacEtmImpl(UmacImpl):
is_etm = True
class Umac128Impl(UmacImpl):
length = 16
umac_cls = umac.Umac128
class Umac128EtmImpl(Umac128Impl):
is_etm = True
class HmacImpl(MacInterface):
hash_cls = None
def verify(self, message: bytes, mac: bytes) -> bool:
got = hmac.HMAC(self.key, message, self.hash_cls).digest()
return secrets.compare_digest(got, mac)
def calculate(self, message: bytes) -> bytes:
got = hmac.HMAC(self.key, message, self.hash_cls).digest()
return got
class HmacSha256Impl(HmacImpl):
length = 32
key_size = 32
hash_cls = hashlib.sha256
is_etm = False
class HmacSha256EtmImpl(HmacSha256Impl):
is_etm = True
class HmacSha512Impl(HmacImpl):
length = 64
key_size = 64
hash_cls = hashlib.sha512
is_etm = False
class HmacSha512EtmImpl(HmacSha512Impl):
is_etm = True
class HmacSha1Impl(HmacImpl):
length = 20
key_size = 20
hash_cls = hashlib.sha1
class HmacSha1EtmImpl(HmacSha1Impl):
is_etm = True
def get_mac_impl(algo: str) -> t.Type["MacInterface"]:
mapping = {
"[email protected]": HmacSha256EtmImpl,
"[email protected]": HmacSha512EtmImpl,
"[email protected]": HmacSha1EtmImpl,
"hmac-sha2-256": HmacSha256Impl,
"hmac-sha2-512": HmacSha512Impl,
"hmac-sha1": HmacSha1Impl,
}
if umac_support:
mapping.update(
{
"[email protected]": UmacImpl,
"[email protected]": UmacEtmImpl,
"[email protected]": Umac128EtmImpl,
}
)
return mapping[algo]