Coverage for tcprocd/client.py: 97.37%

1"""tcprocd client.""" 

2from __future__ import unicode_literals, print_function, absolute_import 

3from tcprocd.protocol import Protocol 

4import socket 

5import select 

6import sys 

7 

8 

9if sys.version_info[0] < 3: 9 ↛ 10line 9 didn't jump to line 10, because the condition on line 9 was never true

10 str_types = (str, unicode) # noqa 

11else: 

12 str_types = (str, bytes) # noqa 

13 

14 

15class SocketShell(object): 

16 """ 

17 A class to connect to a process's thread. 

18 

19 Adds a line buffer for the socket and passes received messages to ``on_receive``. 

20 Messages by the user are passed to ``on_stdin``. 

21 

22 :param client: :class:`tcprocd.client.Client` - The client to use for the connection. 

23 """ 

24 

25 def __init__(self, client): 

26 """Initialize shell.""" 

27 self.client = client 

28 self.sockets = [sys.stdin, client.socket] 

29 self._do_stop = False 

30 

31 def on_stdin_ready(self): 

32 """Read line from stdin and send it.""" 

33 line = sys.stdin.readline().strip() 

34 

35 if line == 'exit': 

36 return True 

37 

38 self.client.protocol.sendline(line + '\n') 

39 

40 def on_socket_ready(self): 

41 """Receive line and write it to stdout.""" 

42 line = self.client.protocol.readline() 

43 

44 if line == 'exit': 

45 return True 

46 

47 sys.stdout.write(line + '\n') 

48 sys.stdout.flush() 

49 

50 def run(self): 

51 """Start waiting for input/output.""" 

52 try: 

53 while not self._do_stop: 

54 ready = select.select(self.sockets, [], [])[0] 

55 

56 for s in ready: 

57 if s == self.client.socket: # message by server 

58 if self.on_socket_ready(): 

59 self._do_stop = True 

60 

61 else: # message by user 

62 if self.on_stdin_ready(): 

63 self._do_stop = True 

64 finally: 

65 self.client.socket.close() 

66 

67 

68class AuthenticationError(Exception): 

69 """Exception raised when authentication fails.""" 

70 

71 pass 

72 

73 

74class ServerError(Exception): 

75 """Exception raised when the server answers with an error.""" 

76 

77 pass 

78 

79 

80class Client(object): 

81 """ 

82 A class to connect to a tcprocd server. 

83 

84 :param server_address: tuple of host and port or the path to the socket file 

85 """ 

86 

87 def __init__(self, server_address): 

88 """Initialize client.""" 

89 self.server_address = server_address 

90 

91 if isinstance(server_address, str_types): 

92 self.is_unix_domain = True 

93 sock_type = socket.AF_UNIX 

94 else: 

95 self.is_unix_domain = False 

96 sock_type = socket.AF_INET 

97 

98 self.socket = socket.socket(sock_type, socket.SOCK_STREAM) 

99 self.protocol = Protocol(self.socket) 

100 self.server_version = None 

101 self.attached_to = None 

102 

103 def connect(self, username=None, password=None, username_callback=None, password_callback=None): 

104 """Connect to the server.""" 

105 if self.is_unix_domain: 

106 try: 

107 SO_PASSCRED = socket.SO_PASSCRED 

108 except AttributeError: 

109 SO_PASSCRED = 16 

110 

111 self.socket.setsockopt(socket.SOL_SOCKET, SO_PASSCRED, 1) 

112 

113 self.socket.connect(self.server_address) 

114 self.server_version = self.protocol.recv_part(3) 

115 

116 answer = self.protocol.recv_part(2) 

117 

118 # TCP connections always require username and password. 

119 # A unix domain socket does not accept an username and 

120 # only requires a password if the connecting user has one. 

121 # TODO: return 'authentication required' and let the caller authenticate on its own 

122 if answer == self.protocol.AUTHENTICATION_REQUIRED: 

123 if not self.is_unix_domain: 

124 

125 if username is None: 

126 username = username_callback() 

127 self.protocol.send_part(2, username) 

128 

129 if password is None: 

130 password = password_callback() 

131 self.protocol.send_part(2, password) 

132 

133 answer = self.protocol.recv_part(2) 

134 

135 if answer != self.protocol.OK: 

136 raise AuthenticationError() 

137 

138 elif answer != self.protocol.OK: 

139 raise ServerError(answer) 

140 

141 def close(self): 

142 """Close the connection.""" 

143 self.socket.close() 

144 

145 def __enter__(self): 

146 """Connect when used as context manager.""" 

147 self.connect() 

148 return self 

149 

150 def __exit__(self, exc_type, exc_val, exc_tb): 

151 """Disconnect afterwards, when used as context manager.""" 

152 self.close() 

153 

154 def list(self): 

155 """List servers.""" 

156 self.protocol.send_part(2, 'list') 

157 answer = self.protocol.recv_part(2) 

158 

159 data = [] 

160 if answer == self.protocol.OK: 

161 data = self.protocol.recv_part(6).split('\n') 

162 

163 return answer, data 

164 

165 def cat(self, name, start=0): 

166 """Get output of given process. 

167 

168 :param name: :class:`str` - Name of the process. 

169 :param start: :class:`int` - Start at this line. (Default: ``0``) 

170 :return: :class:`str` - Multi-line output of the process. 

171 """ 

172 self.protocol.send_part(2, 'cat') 

173 self.protocol.send_part(2, name) 

174 self.protocol.send_part(1, str(start)) 

175 answer = self.protocol.recv_part(2) 

176 data = [] 

177 if answer == self.protocol.OK: 

178 data = self.protocol.recv_part(6).split('\n') 

179 return answer, data 

180 

181 def start(self, name, command, path=''): 

182 """Create a new process with the given ``name`` and ``command``. 

183 

184 :param name: :class:`str` - Name of the process. 

185 :param command: :class:`str` - The command to run the process. 

186 :param path: :class:`str` - The (remote) path to execute the 

187 command in. (Default: ``None``) 

188 :return: :class:`str` - Status message 

189 """ 

190 self.protocol.send_part(2, 'start') 

191 self.protocol.send_part(2, name) 

192 self.protocol.send_part(3, command) 

193 self.protocol.send_part(3, path) 

194 return self.protocol.recv_part(2) 

195 

196 def kill(self, name): 

197 """Kill the given process. 

198 

199 :param name: :class:`str` - Name of the process. 

200 :return: :class:`str` - Status message 

201 """ 

202 self.protocol.send_part(2, 'kill') 

203 self.protocol.send_part(2, name) 

204 return self.protocol.recv_part(2) 

205 

206 def command(self, name, command): 

207 """Write the given command to the given process's stdin. 

208 

209 .. Note: Use ``cat`` to see the process's stdout! 

210 

211 :param name: :class:`str` - Name of the process. 

212 :param command: :class:`str` - The command to send to the process. 

213 :return: :class:`str` - Status message 

214 """ 

215 self.protocol.send_part(2, 'command') 

216 self.protocol.send_part(2, name) 

217 self.protocol.send_part(3, command) 

218 return self.protocol.recv_part(2) 

219 

220 def attach(self, name): 

221 """ 

222 Attach to the given process's shell. 

223 

224 :param name: :class:`str` - Name of the process. 

225 :return: :class:`str` - Status message 

226 """ 

227 self.protocol.send_part(2, 'attach') 

228 self.protocol.send_part(2, name) 

229 return self.protocol.recv_part(2)