Coverage for tcprocd/protocol.py: 100.00%

1"""tcprocd protocol.""" 

2from __future__ import unicode_literals, print_function, absolute_import 

3import logging 

4 

5 

6logger = logging.getLogger(__name__) 

7 

8 

9class ProtocolError(Exception): 

10 """Exception, raised on protocol errors.""" 

11 

12 pass 

13 

14 

15class Disconnect(Exception): 

16 """Exception, raised on disconnect.""" 

17 

18 

19class Protocol(object): 

20 """ 

21 The tcprocd protocol. 

22 

23 :param socket: socket to use for sending/receiving 

24 :param timeout: default timeout between messages 

25 """ 

26 

27 OK = 'OK' 

28 OFFLINE = 'OFFLINE' 

29 EXISTS = 'EXISTS' 

30 UNKNOWN = 'UNKNOWN' 

31 PERMISSION_DENIED = 'PERMISSION_DENIED' 

32 ERROR = 'PROTOCOL_ERROR' 

33 AUTHENTICATION_REQUIRED = 'AUTHENTICATION_REQUIRED' 

34 AUTHENTICATION_ERROR = 'AUTHENTICATION_ERROR' 

35 

36 def __init__(self, socket, timeout=2): 

37 """Initialize protocol.""" 

38 self._timeout = timeout 

39 self.socket = socket 

40 

41 self.set_timeout() 

42 

43 def set_timeout(self, timeout=None): 

44 """Set the given timeout or reset if not given.""" 

45 if timeout is None: 

46 timeout = self._timeout 

47 self.socket.settimeout(timeout) 

48 

49 def _recv(self, length): 

50 received_bytes = self.socket.recv(length) 

51 

52 if not received_bytes: 

53 raise Disconnect 

54 

55 return received_bytes 

56 

57 def recv_part(self, prefix_numbers): 

58 """Receive a string prefixed by the given length and return it.""" 

59 part_len = self._recv(prefix_numbers) 

60 

61 try: 

62 part_len = int(part_len) 

63 except ValueError: 

64 raise ProtocolError('Invalid prefix: "{}"'.format(part_len)) 

65 

66 if part_len == 0: 

67 return '' 

68 return self._recv(part_len).decode() 

69 

70 def readline(self): 

71 """Receive a line and return it.""" 

72 buf = b'' 

73 while b'\n' not in buf: 

74 buf += self._recv(1) 

75 return buf.decode().rstrip('\n') 

76 

77 def sendline(self, line): 

78 """Send a line.""" 

79 if not isinstance(line, bytes): 

80 line = line.encode() 

81 

82 if not line.endswith(b'\n'): 

83 line += b'\n' 

84 self.socket.send(line) 

85 

86 def send_part(self, prefix_numbers, s): 

87 """Send a string prefixed by its length.""" 

88 if not isinstance(s, bytes): 

89 s = s.encode() 

90 

91 prefix = str(len(s)).encode() 

92 

93 if len(prefix) > prefix_numbers: 

94 msg = 'Provided string length ({}) exceeds the maximum of {}.'.format( 

95 len(s), 

96 '9' * prefix_numbers 

97 ) 

98 raise ProtocolError(msg) 

99 

100 while len(prefix) < prefix_numbers: 

101 prefix = b'0' + prefix 

102 

103 part = prefix + s 

104 self.socket.send(part)