Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/ntlm_auth/rc4.py: 34%

1# Copyright: (c) 2018, Jordan Borean (@jborean93) <jborean93@gmail.com> 

2# MIT License (see LICENSE or https://opensource.org/licenses/MIT) 

3 

4import struct 

5 

6# Favour cryptography over our Python implementation as it is a lot faster 

7try: 

8 from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes 

9 from cryptography.hazmat.backends import default_backend 

10 HAS_CRYPTOGRAPHY = True 

11except ImportError: # pragma: no cover 

12 HAS_CRYPTOGRAPHY = False 

13 

14 

15class _CryptographyARC4(object): 

16 

17 def __init__(self, key): 

18 algo = algorithms.ARC4(key) 

19 cipher = Cipher(algo, mode=None, backend=default_backend()) 

20 self._encryptor = cipher.encryptor() 

21 

22 def update(self, value): 

23 return self._encryptor.update(value) 

24 

25 

26class _PythonARC4(object): 

27 state = None 

28 i = 0 

29 j = 0 

30 

31 def __init__(self, key): 

32 # Split up the key into a list 

33 key_bytes = [] 

34 for i in range(len(key)): 

35 key_byte = struct.unpack("B", key[i:i + 1])[0] 

36 key_bytes.append(key_byte) 

37 

38 # Key-scheduling algorithm (KSA) 

39 self.state = [n for n in range(256)] 

40 j = 0 

41 for i in range(256): 

42 j = (j + self.state[i] + key_bytes[i % len(key_bytes)]) % 256 

43 self.state[i], self.state[j] = self.state[j], self.state[i] 

44 

45 def update(self, value): 

46 chars = [] 

47 random_gen = self._random_generator() 

48 for i in range(len(value)): 

49 byte = struct.unpack("B", value[i:i + 1])[0] 

50 updated_byte = byte ^ next(random_gen) 

51 chars.append(updated_byte) 

52 return bytes(bytearray(chars)) 

53 

54 def _random_generator(self): 

55 # Pseudo-Random Generation Algorithm (PRGA) 

56 while True: 

57 self.i = (self.i + 1) % 256 

58 self.j = (self.j + self.state[self.i]) % 256 

59 self.state[self.i], self.state[self.j] = \ 

60 self.state[self.j], self.state[self.i] 

61 yield self.state[(self.state[self.i] + self.state[self.j]) % 256] 

62 

63 

64if HAS_CRYPTOGRAPHY: 

65 ARC4 = _CryptographyARC4 

66else: # pragma: no cover 

67 ARC4 = _PythonARC4