blob: 17faff040f14d0a38d88c1f160a7c4de92e5c3e2 [file] [log] [blame]
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001#!/usr/bin/python -u
2# -*- coding: utf-8 -*-
3#
4# Copyright 2015 The Chromium OS Authors. All rights reserved.
5# Use of this source code is governed by a BSD-style license that can be
6# found in the LICENSE file.
7
Wei-Ning Huang7d029b12015-03-06 10:32:15 +08008import argparse
Wei-Ning Huangb05cde32015-08-01 09:48:41 +08009import contextlib
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080010import fcntl
Wei-Ning Huangb05cde32015-08-01 09:48:41 +080011import hashlib
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080012import json
13import logging
14import os
15import Queue
Wei-Ning Huang829e0c82015-05-26 14:37:23 +080016import re
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080017import select
Wei-Ning Huanga301f572015-06-03 17:34:21 +080018import signal
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080019import socket
Wei-Ning Huangb05cde32015-08-01 09:48:41 +080020import struct
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080021import subprocess
22import sys
Moja Hsuc9ecc8b2015-07-13 11:39:17 +080023import termios
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080024import threading
25import time
Wei-Ning Huangb05cde32015-08-01 09:48:41 +080026import urllib
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080027import uuid
28
Wei-Ning Huang2132de32015-04-13 17:24:38 +080029import jsonrpclib
30from jsonrpclib.SimpleJSONRPCServer import SimpleJSONRPCServer
31
32
33_GHOST_RPC_PORT = 4499
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080034
35_OVERLORD_PORT = 4455
36_OVERLORD_LAN_DISCOVERY_PORT = 4456
Wei-Ning Huangb05cde32015-08-01 09:48:41 +080037_OVERLORD_HTTP_PORT = 9000
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080038
39_BUFSIZE = 8192
40_RETRY_INTERVAL = 2
41_SEPARATOR = '\r\n'
42_PING_TIMEOUT = 3
43_PING_INTERVAL = 5
44_REQUEST_TIMEOUT_SECS = 60
45_SHELL = os.getenv('SHELL', '/bin/bash')
Wei-Ning Huang2132de32015-04-13 17:24:38 +080046_DEFAULT_BIND_ADDRESS = '0.0.0.0'
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080047
Moja Hsuc9ecc8b2015-07-13 11:39:17 +080048_CONTROL_START = 128
49_CONTROL_END = 129
50
Wei-Ning Huanga301f572015-06-03 17:34:21 +080051_BLOCK_SIZE = 4096
52
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080053RESPONSE_SUCCESS = 'success'
54RESPONSE_FAILED = 'failed'
55
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080056class PingTimeoutError(Exception):
57 pass
58
59
60class RequestError(Exception):
61 pass
62
63
64class Ghost(object):
65 """Ghost implements the client protocol of Overlord.
66
67 Ghost provide terminal/shell/logcat functionality and manages the client
68 side connectivity.
69 """
Wei-Ning Huanga301f572015-06-03 17:34:21 +080070 NONE, AGENT, TERMINAL, SHELL, LOGCAT, FILE = range(6)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080071
72 MODE_NAME = {
73 NONE: 'NONE',
74 AGENT: 'Agent',
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +080075 TERMINAL: 'Terminal',
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080076 SHELL: 'Shell',
Wei-Ning Huanga301f572015-06-03 17:34:21 +080077 LOGCAT: 'Logcat',
78 FILE: 'File'
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080079 }
80
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +080081 RANDOM_MID = '##random_mid##'
82
Wei-Ning Huanga301f572015-06-03 17:34:21 +080083 def __init__(self, overlord_addrs, mode=AGENT, mid=None, sid=None, bid=None,
84 command=None, file_op=None):
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080085 """Constructor.
86
87 Args:
88 overlord_addrs: a list of possible address of overlord.
89 mode: client mode, either AGENT, SHELL or LOGCAT
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +080090 mid: a str to set for machine ID. If mid equals Ghost.RANDOM_MID, machine
91 id is randomly generated.
Wei-Ning Huanga301f572015-06-03 17:34:21 +080092 sid: session ID. If the connection is requested by overlord, sid should
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080093 be set to the corresponding session id assigned by overlord.
Wei-Ning Huanga301f572015-06-03 17:34:21 +080094 bid: browser ID. Identifies the browser which started the session.
95 command: the command to execute when we are in SHELL mode.
Wei-Ning Huange2981862015-08-03 15:03:08 +080096 file_op: a tuple (action, filepath, pid). action is either 'download' or
97 'upload'. pid is the pid of the target shell, used to determine where
98 the current working is and thus where to upload to.
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080099 """
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800100 assert mode in [Ghost.AGENT, Ghost.TERMINAL, Ghost.SHELL, Ghost.FILE]
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800101 if mode == Ghost.SHELL:
102 assert command is not None
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800103 if mode == Ghost.FILE:
104 assert file_op is not None
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800105
106 self._overlord_addrs = overlord_addrs
Wei-Ning Huangad330c52015-03-12 20:34:18 +0800107 self._connected_addr = None
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800108 self._mode = mode
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +0800109 self._mid = mid
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800110 self._sock = None
111 self._machine_id = self.GetMachineID()
Wei-Ning Huangfed95862015-08-07 03:17:11 +0800112 self._session_id = sid if sid is not None else str(uuid.uuid4())
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800113 self._browser_id = bid
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800114 self._properties = {}
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800115 self._shell_command = command
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800116 self._file_op = file_op
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800117 self._buf = ''
118 self._requests = {}
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800119 self._reset = threading.Event()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800120 self._last_ping = 0
121 self._queue = Queue.Queue()
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800122 self._download_queue = Queue.Queue()
Wei-Ning Huange2981862015-08-03 15:03:08 +0800123 self._ttyname_to_bid = {}
124 self._terminal_sid_to_pid = {}
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800125
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800126 def SetIgnoreChild(self, status):
127 # Only ignore child for Agent since only it could spawn child Ghost.
128 if self._mode == Ghost.AGENT:
129 signal.signal(signal.SIGCHLD,
130 signal.SIG_IGN if status else signal.SIG_DFL)
131
132 def GetFileSha1(self, filename):
133 with open(filename, 'r') as f:
134 return hashlib.sha1(f.read()).hexdigest()
135
136 def Upgrade(self):
137 logging.info('Upgrade: initiating upgrade sequence...')
138
139 scriptpath = os.path.abspath(sys.argv[0])
140 url = 'http://%s:%d/upgrade/ghost.py' % (
141 self._connected_addr[0], _OVERLORD_HTTP_PORT)
142
143 # Download sha1sum for ghost.py for verification
144 try:
145 with contextlib.closing(urllib.urlopen(url + '.sha1')) as f:
146 if f.getcode() != 200:
147 raise RuntimeError('HTTP status %d' % f.getcode())
148 sha1sum = f.read().strip()
149 except Exception:
150 logging.error('Upgrade: failed to download sha1sum file, abort')
151 return
152
153 if self.GetFileSha1(scriptpath) == sha1sum:
154 logging.info('Upgrade: ghost is already up-to-date, skipping upgrade')
155 return
156
157 # Download upgrade version of ghost.py
158 try:
159 with contextlib.closing(urllib.urlopen(url)) as f:
160 if f.getcode() != 200:
161 raise RuntimeError('HTTP status %d' % f.getcode())
162 data = f.read()
163 except Exception:
164 logging.error('Upgrade: failed to download upgrade, abort')
165 return
166
167 # Compare SHA1 sum
168 if hashlib.sha1(data).hexdigest() != sha1sum:
169 logging.error('Upgrade: sha1sum mismatch, abort')
170 return
171
172 python = os.readlink('/proc/self/exe')
173 try:
174 with open(scriptpath, 'w') as f:
175 f.write(data)
176 except Exception:
177 logging.error('Upgrade: failed to write upgrade onto disk, abort')
178 return
179
180 logging.info('Upgrade: restarting ghost...')
181 self.CloseSockets()
182 self.SetIgnoreChild(False)
183 os.execve(python, [python, scriptpath] + sys.argv[1:], os.environ)
184
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800185 def LoadPropertiesFromFile(self, filename):
186 try:
187 with open(filename, 'r') as f:
188 self._properties = json.loads(f.read())
189 except Exception as e:
190 logging.exception('LoadPropertiesFromFile: ' + str(e))
191
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800192 def CloseSockets(self):
193 # Close sockets opened by parent process, since we don't use it anymore.
194 for fd in os.listdir('/proc/self/fd/'):
195 try:
196 real_fd = os.readlink('/proc/self/fd/%s' % fd)
197 if real_fd.startswith('socket'):
198 os.close(int(fd))
199 except Exception:
200 pass
201
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800202 def SpawnGhost(self, mode, sid=None, bid=None, command=None, file_op=None):
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800203 """Spawn a child ghost with specific mode.
204
205 Returns:
206 The spawned child process pid.
207 """
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800208 # Restore the default signal hanlder, so our child won't have problems.
209 self.SetIgnoreChild(False)
210
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800211 pid = os.fork()
212 if pid == 0:
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800213 self.CloseSockets()
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800214 g = Ghost([self._connected_addr], mode, Ghost.RANDOM_MID, sid, bid,
215 command, file_op)
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800216 g.Start()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800217 sys.exit(0)
218 else:
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800219 self.SetIgnoreChild(True)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800220 return pid
221
222 def Timestamp(self):
223 return int(time.time())
224
225 def GetGateWayIP(self):
226 with open('/proc/net/route', 'r') as f:
227 lines = f.readlines()
228
229 ips = []
230 for line in lines:
231 parts = line.split('\t')
232 if parts[2] == '00000000':
233 continue
234
235 try:
236 h = parts[2].decode('hex')
237 ips.append('%d.%d.%d.%d' % tuple(ord(x) for x in reversed(h)))
238 except TypeError:
239 pass
240
241 return ips
242
Wei-Ning Huang829e0c82015-05-26 14:37:23 +0800243 def GetShopfloorIP(self):
244 try:
245 import factory_common # pylint: disable=W0612
246 from cros.factory.test import shopfloor
247
248 url = shopfloor.get_server_url()
249 match = re.match(r'^https?://(.*):.*$', url)
250 if match:
251 return [match.group(1)]
252 except Exception:
253 pass
254 return []
255
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800256 def GetMachineID(self):
257 """Generates machine-dependent ID string for a machine.
258 There are many ways to generate a machine ID:
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800259 1. factory device_id
260 2. factory device-data
Wei-Ning Huang1d7603b2015-07-03 17:38:56 +0800261 3. /sys/class/dmi/id/product_uuid (only available on intel machines)
262 4. MAC address
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800263 We follow the listed order to generate machine ID, and fallback to the next
264 alternative if the previous doesn't work.
265 """
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +0800266 if self._mid == Ghost.RANDOM_MID:
Wei-Ning Huangaed90452015-03-23 17:50:21 +0800267 return str(uuid.uuid4())
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +0800268 elif self._mid:
269 return self._mid
Wei-Ning Huangaed90452015-03-23 17:50:21 +0800270
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800271 # Try factory device id
272 try:
273 import factory_common # pylint: disable=W0612
274 from cros.factory.test import event_log
275 with open(event_log.DEVICE_ID_PATH) as f:
276 return f.read().strip()
277 except Exception:
278 pass
279
Wei-Ning Huang1d7603b2015-07-03 17:38:56 +0800280 # Try factory device data
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800281 try:
282 p = subprocess.Popen('factory device-data | grep mlb_serial_number | '
283 'cut -d " " -f 2', stdout=subprocess.PIPE,
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800284 stderr=subprocess.PIPE, shell=True)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800285 stdout, _ = p.communicate()
286 if stdout == '':
Wei-Ning Huang1d7603b2015-07-03 17:38:56 +0800287 raise RuntimeError('empty mlb number')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800288 return stdout.strip()
289 except Exception:
290 pass
291
Wei-Ning Huang1d7603b2015-07-03 17:38:56 +0800292 # Try DMI product UUID
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800293 try:
294 with open('/sys/class/dmi/id/product_uuid', 'r') as f:
295 return f.read().strip()
296 except Exception:
297 pass
298
Wei-Ning Huang1d7603b2015-07-03 17:38:56 +0800299 # Use MAC address if non is available
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800300 try:
301 macs = []
302 ifaces = sorted(os.listdir('/sys/class/net'))
303 for iface in ifaces:
304 if iface == 'lo':
305 continue
306
307 with open('/sys/class/net/%s/address' % iface, 'r') as f:
308 macs.append(f.read().strip())
309
310 return ';'.join(macs)
311 except Exception:
312 pass
313
Moja Hsuc9ecc8b2015-07-13 11:39:17 +0800314 raise RuntimeError('can\'t generate machine ID')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800315
316 def Reset(self):
317 """Reset state and clear request handlers."""
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800318 self._reset.clear()
Moja Hsuc9ecc8b2015-07-13 11:39:17 +0800319 self._buf = ''
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800320 self._last_ping = 0
321 self._requests = {}
322
323 def SendMessage(self, msg):
324 """Serialize the message and send it through the socket."""
325 self._sock.send(json.dumps(msg) + _SEPARATOR)
326
327 def SendRequest(self, name, args, handler=None,
328 timeout=_REQUEST_TIMEOUT_SECS):
329 if handler and not callable(handler):
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800330 raise RequestError('Invalid request handler for msg "%s"' % name)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800331
332 rid = str(uuid.uuid4())
333 msg = {'rid': rid, 'timeout': timeout, 'name': name, 'params': args}
Wei-Ning Huange2981862015-08-03 15:03:08 +0800334 if timeout >= 0:
335 self._requests[rid] = [self.Timestamp(), timeout, handler]
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800336 self.SendMessage(msg)
337
338 def SendResponse(self, omsg, status, params=None):
339 msg = {'rid': omsg['rid'], 'response': status, 'params': params}
340 self.SendMessage(msg)
341
Moja Hsuc9ecc8b2015-07-13 11:39:17 +0800342 def HandlePTYControl(self, fd, control_string):
343 msg = json.loads(control_string)
344 command = msg['command']
345 params = msg['params']
346 if command == 'resize':
347 # some error happened on websocket
348 if len(params) != 2:
349 return
350 winsize = struct.pack('HHHH', params[0], params[1], 0, 0)
351 fcntl.ioctl(fd, termios.TIOCSWINSZ, winsize)
352 else:
353 logging.warn('Invalid request command "%s"', command)
354
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800355 def SpawnPTYServer(self, _):
356 """Spawn a PTY server and forward I/O to the TCP socket."""
357 logging.info('SpawnPTYServer: started')
358
359 pid, fd = os.forkpty()
360 if pid == 0:
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800361 # Register the mapping of browser_id and ttyname
362 ttyname = os.readlink('/proc/%d/fd/0' % os.getpid())
363 try:
364 server = GhostRPCServer()
365 server.RegisterTTY(self._browser_id, ttyname)
Wei-Ning Huange2981862015-08-03 15:03:08 +0800366 server.RegisterSession(self._session_id, os.getpid())
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800367 except Exception:
368 # If ghost is launched without RPC server, the call will fail but we
369 # can ignore it.
370 pass
371
372 # The directory that contains the current running ghost script
373 script_dir = os.path.dirname(os.path.abspath(sys.argv[0]))
374
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800375 env = os.environ.copy()
376 env['USER'] = os.getenv('USER', 'root')
377 env['HOME'] = os.getenv('HOME', '/root')
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800378 env['PATH'] = os.getenv('PATH') + ':%s' % script_dir
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800379 os.chdir(env['HOME'])
380 os.execve(_SHELL, [_SHELL], env)
381 else:
382 try:
Moja Hsuc9ecc8b2015-07-13 11:39:17 +0800383 control_state = None
384 control_string = ''
385 write_buffer = ''
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800386 while True:
387 rd, _, _ = select.select([self._sock, fd], [], [])
388
389 if fd in rd:
390 self._sock.send(os.read(fd, _BUFSIZE))
391
392 if self._sock in rd:
393 ret = self._sock.recv(_BUFSIZE)
394 if len(ret) == 0:
Moja Hsuc9ecc8b2015-07-13 11:39:17 +0800395 raise RuntimeError('socket closed')
396 while ret:
397 if control_state:
398 if chr(_CONTROL_END) in ret:
399 index = ret.index(chr(_CONTROL_END))
400 control_string += ret[:index]
401 self.HandlePTYControl(fd, control_string)
402 control_state = None
403 control_string = ''
404 ret = ret[index+1:]
405 else:
406 control_string += ret
407 ret = ''
408 else:
409 if chr(_CONTROL_START) in ret:
410 control_state = _CONTROL_START
411 index = ret.index(chr(_CONTROL_START))
412 write_buffer += ret[:index]
413 ret = ret[index+1:]
414 else:
415 write_buffer += ret
416 ret = ''
417 if write_buffer:
418 os.write(fd, write_buffer)
419 write_buffer = ''
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800420 except (OSError, socket.error, RuntimeError):
421 self._sock.close()
422 logging.info('SpawnPTYServer: terminated')
423 sys.exit(0)
424
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800425 def SpawnShellServer(self, _):
426 """Spawn a shell server and forward input/output from/to the TCP socket."""
427 logging.info('SpawnShellServer: started')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800428
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800429 p = subprocess.Popen(self._shell_command, stdin=subprocess.PIPE,
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800430 stdout=subprocess.PIPE, stderr=subprocess.PIPE,
431 shell=True)
432
433 def make_non_block(fd):
434 fl = fcntl.fcntl(fd, fcntl.F_GETFL)
435 fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
436
437 make_non_block(p.stdout)
438 make_non_block(p.stderr)
439
440 try:
441 while True:
442 rd, _, _ = select.select([p.stdout, p.stderr, self._sock], [], [])
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800443 p.poll()
444
445 if p.returncode != None:
Moja Hsuc9ecc8b2015-07-13 11:39:17 +0800446 raise RuntimeError('process complete')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800447
448 if p.stdout in rd:
449 self._sock.send(p.stdout.read(_BUFSIZE))
450
451 if p.stderr in rd:
452 self._sock.send(p.stderr.read(_BUFSIZE))
453
454 if self._sock in rd:
455 ret = self._sock.recv(_BUFSIZE)
456 if len(ret) == 0:
Moja Hsuc9ecc8b2015-07-13 11:39:17 +0800457 raise RuntimeError('socket closed')
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800458 p.stdin.write(ret)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800459 except (OSError, socket.error, RuntimeError):
460 self._sock.close()
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800461 logging.info('SpawnShellServer: terminated')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800462 sys.exit(0)
463
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800464 def InitiateFileOperation(self, _):
465 if self._file_op[0] == 'download':
466 size = os.stat(self._file_op[1]).st_size
467 self.SendRequest('request_to_download',
468 {'bid': self._browser_id,
469 'filename': os.path.basename(self._file_op[1]),
470 'size': size})
Wei-Ning Huange2981862015-08-03 15:03:08 +0800471 elif self._file_op[0] == 'upload':
472 self.SendRequest('clear_to_upload', {}, timeout=-1)
473 self.StartUploadServer()
474 else:
475 logging.error('InitiateFileOperation: unknown file operation, ignored')
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800476
477 def StartDownloadServer(self):
478 logging.info('StartDownloadServer: started')
479
480 try:
481 with open(self._file_op[1], 'rb') as f:
482 while True:
483 data = f.read(_BLOCK_SIZE)
484 if len(data) == 0:
485 break
486 self._sock.send(data)
487 except Exception as e:
488 logging.error('StartDownloadServer: %s', e)
489 finally:
490 self._sock.close()
491
492 logging.info('StartDownloadServer: terminated')
493 sys.exit(0)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800494
Wei-Ning Huange2981862015-08-03 15:03:08 +0800495 def StartUploadServer(self):
496 logging.info('StartUploadServer: started')
497
498 try:
499 target_dir = os.getenv('HOME', '/tmp')
500
501 # Get the client's working dir, which is our target upload dir
502 if self._file_op[2]:
503 target_dir = os.readlink('/proc/%d/cwd' % self._file_op[2])
504
505 self._sock.setblocking(False)
506 with open(os.path.join(target_dir, self._file_op[1]), 'wb') as f:
507 while True:
508 rd, _, _ = select.select([self._sock], [], [])
509 if self._sock in rd:
510 buf = self._sock.recv(_BLOCK_SIZE)
511 if len(buf) == 0:
512 break
513 f.write(buf)
514 except socket.error as e:
515 logging.error('StartUploadServer: socket error: %s', e)
516 except Exception as e:
517 logging.error('StartUploadServer: %s', e)
518 finally:
519 self._sock.close()
520
521 logging.info('StartUploadServer: terminated')
522 sys.exit(0)
523
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800524 def Ping(self):
525 def timeout_handler(x):
526 if x is None:
527 raise PingTimeoutError
528
529 self._last_ping = self.Timestamp()
530 self.SendRequest('ping', {}, timeout_handler, 5)
531
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800532 def HandleRequest(self, msg):
Wei-Ning Huange2981862015-08-03 15:03:08 +0800533 command = msg['name']
534 params = msg['params']
535
536 if command == 'upgrade':
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800537 self.Upgrade()
Wei-Ning Huange2981862015-08-03 15:03:08 +0800538 elif command == 'terminal':
539 self.SpawnGhost(self.TERMINAL, params['sid'], bid=params['bid'])
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800540 self.SendResponse(msg, RESPONSE_SUCCESS)
Wei-Ning Huange2981862015-08-03 15:03:08 +0800541 elif command == 'shell':
542 self.SpawnGhost(self.SHELL, params['sid'], command=params['command'])
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800543 self.SendResponse(msg, RESPONSE_SUCCESS)
Wei-Ning Huange2981862015-08-03 15:03:08 +0800544 elif command == 'file_download':
545 self.SpawnGhost(self.FILE, params['sid'],
546 file_op=('download', params['filename'], None))
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800547 self.SendResponse(msg, RESPONSE_SUCCESS)
Wei-Ning Huange2981862015-08-03 15:03:08 +0800548 elif command == 'clear_to_download':
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800549 self.StartDownloadServer()
Wei-Ning Huange2981862015-08-03 15:03:08 +0800550 elif command == 'file_upload':
551 pid = self._terminal_sid_to_pid.get(params['terminal_sid'], None)
552 self.SpawnGhost(self.FILE, params['sid'],
553 file_op=('upload', params['filename'], pid))
554 self.SendResponse(msg, RESPONSE_SUCCESS)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800555
556 def HandleResponse(self, response):
557 rid = str(response['rid'])
558 if rid in self._requests:
559 handler = self._requests[rid][2]
560 del self._requests[rid]
561 if callable(handler):
562 handler(response)
563 else:
564 print(response, self._requests.keys())
565 logging.warning('Recvied unsolicited response, ignored')
566
567 def ParseMessage(self):
568 msgs_json = self._buf.split(_SEPARATOR)
569 self._buf = msgs_json.pop()
570
571 for msg_json in msgs_json:
572 try:
573 msg = json.loads(msg_json)
574 except ValueError:
575 # Ignore mal-formed message.
576 continue
577
578 if 'name' in msg:
579 self.HandleRequest(msg)
580 elif 'response' in msg:
581 self.HandleResponse(msg)
582 else: # Ingnore mal-formed message.
583 pass
584
585 def ScanForTimeoutRequests(self):
586 for rid in self._requests.keys()[:]:
587 request_time, timeout, handler = self._requests[rid]
588 if self.Timestamp() - request_time > timeout:
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800589 if callable(handler):
590 handler(None)
591 else:
592 logging.error('Request %s timeout', rid)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800593 del self._requests[rid]
594
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800595 def InitiateDownload(self):
596 ttyname, filename = self._download_queue.get()
Wei-Ning Huange2981862015-08-03 15:03:08 +0800597 bid = self._ttyname_to_bid[ttyname]
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800598 self.SpawnGhost(self.FILE, bid=bid, file_op=('download', filename))
599
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800600 def Listen(self):
601 try:
602 while True:
603 rds, _, _ = select.select([self._sock], [], [], _PING_INTERVAL / 2)
604
605 if self._sock in rds:
606 self._buf += self._sock.recv(_BUFSIZE)
607 self.ParseMessage()
608
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800609 if (self._mode == self.AGENT and
610 self.Timestamp() - self._last_ping > _PING_INTERVAL):
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800611 self.Ping()
612 self.ScanForTimeoutRequests()
613
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800614 if not self._download_queue.empty():
615 self.InitiateDownload()
616
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800617 if self._reset.is_set():
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800618 self.Reset()
619 break
620 except socket.error:
621 raise RuntimeError('Connection dropped')
622 except PingTimeoutError:
623 raise RuntimeError('Connection timeout')
624 finally:
625 self._sock.close()
626
627 self._queue.put('resume')
628
629 if self._mode != Ghost.AGENT:
630 sys.exit(1)
631
632 def Register(self):
633 non_local = {}
634 for addr in self._overlord_addrs:
635 non_local['addr'] = addr
636 def registered(response):
637 if response is None:
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800638 self._reset.set()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800639 raise RuntimeError('Register request timeout')
640 logging.info('Registered with Overlord at %s:%d', *non_local['addr'])
Moja Hsuc9ecc8b2015-07-13 11:39:17 +0800641 self._queue.put('pause', True)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800642
643 try:
644 logging.info('Trying %s:%d ...', *addr)
645 self.Reset()
646 self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
647 self._sock.settimeout(_PING_TIMEOUT)
648 self._sock.connect(addr)
649
650 logging.info('Connection established, registering...')
651 handler = {
652 Ghost.AGENT: registered,
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800653 Ghost.TERMINAL: self.SpawnPTYServer,
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800654 Ghost.SHELL: self.SpawnShellServer,
655 Ghost.FILE: self.InitiateFileOperation,
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800656 }[self._mode]
657
658 # Machine ID may change if MAC address is used (USB-ethernet dongle
659 # plugged/unplugged)
660 self._machine_id = self.GetMachineID()
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800661 self.SendRequest('register',
662 {'mode': self._mode, 'mid': self._machine_id,
Wei-Ning Huangfed95862015-08-07 03:17:11 +0800663 'sid': self._session_id,
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800664 'properties': self._properties}, handler)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800665 except socket.error:
666 pass
667 else:
668 self._sock.settimeout(None)
Wei-Ning Huangad330c52015-03-12 20:34:18 +0800669 self._connected_addr = addr
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800670 self.Listen()
671
Moja Hsuc9ecc8b2015-07-13 11:39:17 +0800672 raise RuntimeError('Cannot connect to any server')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800673
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800674 def Reconnect(self):
675 logging.info('Received reconnect request from RPC server, reconnecting...')
676 self._reset.set()
677
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800678 def AddToDownloadQueue(self, ttyname, filename):
679 self._download_queue.put((ttyname, filename))
680
681 def RegisterTTY(self, browser_id, ttyname):
Wei-Ning Huange2981862015-08-03 15:03:08 +0800682 self._ttyname_to_bid[ttyname] = browser_id
683
684 def RegisterSession(self, session_id, process_id):
685 self._terminal_sid_to_pid[session_id] = process_id
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800686
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800687 def StartLanDiscovery(self):
688 """Start to listen to LAN discovery packet at
689 _OVERLORD_LAN_DISCOVERY_PORT."""
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800690
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800691 def thread_func():
692 s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
693 s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
694 s.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800695 try:
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800696 s.bind(('0.0.0.0', _OVERLORD_LAN_DISCOVERY_PORT))
697 except socket.error as e:
Moja Hsuc9ecc8b2015-07-13 11:39:17 +0800698 logging.error('LAN discovery: %s, abort', e)
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800699 return
700
701 logging.info('LAN Discovery: started')
702 while True:
703 rd, _, _ = select.select([s], [], [], 1)
704
705 if s in rd:
706 data, source_addr = s.recvfrom(_BUFSIZE)
707 parts = data.split()
708 if parts[0] == 'OVERLORD':
709 ip, port = parts[1].split(':')
710 if not ip:
711 ip = source_addr[0]
712 self._queue.put((ip, int(port)), True)
713
714 try:
715 obj = self._queue.get(False)
716 except Queue.Empty:
717 pass
718 else:
719 if type(obj) is not str:
720 self._queue.put(obj)
721 elif obj == 'pause':
722 logging.info('LAN Discovery: paused')
723 while obj != 'resume':
724 obj = self._queue.get(True)
725 logging.info('LAN Discovery: resumed')
726
727 t = threading.Thread(target=thread_func)
728 t.daemon = True
729 t.start()
730
731 def StartRPCServer(self):
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800732 logging.info("RPC Server: started")
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800733 rpc_server = SimpleJSONRPCServer((_DEFAULT_BIND_ADDRESS, _GHOST_RPC_PORT),
734 logRequests=False)
735 rpc_server.register_function(self.Reconnect, 'Reconnect')
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800736 rpc_server.register_function(self.RegisterTTY, 'RegisterTTY')
Wei-Ning Huange2981862015-08-03 15:03:08 +0800737 rpc_server.register_function(self.RegisterSession, 'RegisterSession')
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800738 rpc_server.register_function(self.AddToDownloadQueue, 'AddToDownloadQueue')
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800739 t = threading.Thread(target=rpc_server.serve_forever)
740 t.daemon = True
741 t.start()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800742
Wei-Ning Huang829e0c82015-05-26 14:37:23 +0800743 def ScanServer(self):
744 for meth in [self.GetGateWayIP, self.GetShopfloorIP]:
745 for addr in [(x, _OVERLORD_PORT) for x in meth()]:
746 if addr not in self._overlord_addrs:
747 self._overlord_addrs.append(addr)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800748
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800749 def Start(self, lan_disc=False, rpc_server=False):
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800750 logging.info('%s started', self.MODE_NAME[self._mode])
751 logging.info('MID: %s', self._machine_id)
Wei-Ning Huangfed95862015-08-07 03:17:11 +0800752 logging.info('SID: %s', self._session_id)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800753
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800754 # We don't care about child process's return code, not wait is needed. This
755 # is used to prevent zombie process from lingering in the system.
756 self.SetIgnoreChild(True)
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800757
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800758 if lan_disc:
759 self.StartLanDiscovery()
760
761 if rpc_server:
762 self.StartRPCServer()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800763
764 try:
765 while True:
766 try:
767 addr = self._queue.get(False)
768 except Queue.Empty:
769 pass
770 else:
771 if type(addr) == tuple and addr not in self._overlord_addrs:
772 logging.info('LAN Discovery: got overlord address %s:%d', *addr)
773 self._overlord_addrs.append(addr)
774
775 try:
Wei-Ning Huang829e0c82015-05-26 14:37:23 +0800776 self.ScanServer()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800777 self.Register()
778 except Exception as e:
779 logging.info(str(e) + ', retrying in %ds' % _RETRY_INTERVAL)
780 time.sleep(_RETRY_INTERVAL)
781
782 self.Reset()
783 except KeyboardInterrupt:
784 logging.error('Received keyboard interrupt, quit')
785 sys.exit(0)
786
787
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800788def GhostRPCServer():
789 return jsonrpclib.Server('http://localhost:%d' % _GHOST_RPC_PORT)
790
791
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800792def DownloadFile(filename):
793 filepath = os.path.abspath(filename)
794 if not os.path.exists(filepath):
795 logging.error("file `%s' does not exist", filename)
796 sys.exit(1)
797
798 # Check if we actually have permission to read the file
799 if not os.access(filepath, os.R_OK):
800 logging.error("can not open %s for reading", filepath)
801 sys.exit(1)
802
803 server = GhostRPCServer()
804 server.AddToDownloadQueue(os.ttyname(0), filepath)
805 sys.exit(0)
806
807
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800808def main():
809 logger = logging.getLogger()
810 logger.setLevel(logging.INFO)
811
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800812 parser = argparse.ArgumentParser()
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +0800813 parser.add_argument('--mid', metavar='MID', dest='mid', action='store',
814 default=None, help='use MID as machine ID')
815 parser.add_argument('--rand-mid', dest='mid', action='store_const',
816 const=Ghost.RANDOM_MID, help='use random machine ID')
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800817 parser.add_argument('--no-lan-disc', dest='lan_disc', action='store_false',
818 default=True, help='disable LAN discovery')
819 parser.add_argument('--no-rpc-server', dest='rpc_server',
820 action='store_false', default=True,
821 help='disable RPC server')
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800822 parser.add_argument('--prop-file', metavar='PROP_FILE', dest="prop_file",
823 type=str, default=None,
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800824 help='file containing the JSON representation of client '
825 'properties')
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800826 parser.add_argument('--download', metavar='FILE', dest='download', type=str,
827 default=None, help='file to download')
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800828 parser.add_argument('overlord_ip', metavar='OVERLORD_IP', type=str,
829 nargs='*', help='overlord server address')
830 args = parser.parse_args()
831
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800832 if args.download:
833 DownloadFile(args.download)
834
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800835 addrs = [('localhost', _OVERLORD_PORT)]
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800836 addrs += [(x, _OVERLORD_PORT) for x in args.overlord_ip]
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800837
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +0800838 g = Ghost(addrs, Ghost.AGENT, args.mid)
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800839 if args.prop_file:
840 g.LoadPropertiesFromFile(args.prop_file)
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800841 g.Start(args.lan_disc, args.rpc_server)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800842
843
844if __name__ == '__main__':
845 main()