blob: a588de8a4a034d719a12064a6c50977baab107eb [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
Joel Kitching22b89042015-08-06 18:23:29 +08008from __future__ import print_function
9
Wei-Ning Huang7d029b12015-03-06 10:32:15 +080010import argparse
Wei-Ning Huangb05cde32015-08-01 09:48:41 +080011import contextlib
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080012import fcntl
Wei-Ning Huangb05cde32015-08-01 09:48:41 +080013import hashlib
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080014import json
15import logging
16import os
17import Queue
Wei-Ning Huang829e0c82015-05-26 14:37:23 +080018import re
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080019import select
Wei-Ning Huanga301f572015-06-03 17:34:21 +080020import signal
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080021import socket
Wei-Ning Huangb05cde32015-08-01 09:48:41 +080022import struct
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080023import subprocess
24import sys
Moja Hsuc9ecc8b2015-07-13 11:39:17 +080025import termios
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080026import threading
27import time
Joel Kitching22b89042015-08-06 18:23:29 +080028import traceback
Wei-Ning Huang39169902015-09-19 06:00:23 +080029import tty
Wei-Ning Huangb05cde32015-08-01 09:48:41 +080030import urllib
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080031import uuid
32
Wei-Ning Huang2132de32015-04-13 17:24:38 +080033import jsonrpclib
34from jsonrpclib.SimpleJSONRPCServer import SimpleJSONRPCServer
35
36
37_GHOST_RPC_PORT = 4499
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080038
39_OVERLORD_PORT = 4455
40_OVERLORD_LAN_DISCOVERY_PORT = 4456
Wei-Ning Huangb05cde32015-08-01 09:48:41 +080041_OVERLORD_HTTP_PORT = 9000
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080042
43_BUFSIZE = 8192
44_RETRY_INTERVAL = 2
45_SEPARATOR = '\r\n'
46_PING_TIMEOUT = 3
47_PING_INTERVAL = 5
48_REQUEST_TIMEOUT_SECS = 60
49_SHELL = os.getenv('SHELL', '/bin/bash')
Wei-Ning Huang2132de32015-04-13 17:24:38 +080050_DEFAULT_BIND_ADDRESS = '0.0.0.0'
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080051
Moja Hsuc9ecc8b2015-07-13 11:39:17 +080052_CONTROL_START = 128
53_CONTROL_END = 129
54
Wei-Ning Huanga301f572015-06-03 17:34:21 +080055_BLOCK_SIZE = 4096
56
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080057RESPONSE_SUCCESS = 'success'
58RESPONSE_FAILED = 'failed'
59
Joel Kitching22b89042015-08-06 18:23:29 +080060
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080061class PingTimeoutError(Exception):
62 pass
63
64
65class RequestError(Exception):
66 pass
67
68
Joel Kitching22b89042015-08-06 18:23:29 +080069class SSHPortForwarder(object):
70 """Create and maintain an SSH port forwarding connection.
71
72 This is meant to be a standalone class to maintain an SSH port forwarding
73 connection to a given server. It provides a fail/retry mechanism, and also
74 can report its current connection status.
75 """
76 _FAILED_STR = 'port forwarding failed'
77 _DEFAULT_CONNECT_TIMEOUT = 10
78 _DEFAULT_ALIVE_INTERVAL = 10
79 _DEFAULT_DISCONNECT_WAIT = 1
80 _DEFAULT_RETRIES = 5
81 _DEFAULT_EXP_FACTOR = 1
82 _DEBUG_INTERVAL = 2
83
84 CONNECTING = 1
85 INITIALIZED = 2
86 FAILED = 4
87
88 REMOTE = 1
89 LOCAL = 2
90
91 @classmethod
92 def ToRemote(cls, *args, **kwargs):
93 """Calls contructor with forward_to=REMOTE."""
94 return cls(*args, forward_to=cls.REMOTE, **kwargs)
95
96 @classmethod
97 def ToLocal(cls, *args, **kwargs):
98 """Calls contructor with forward_to=LOCAL."""
99 return cls(*args, forward_to=cls.LOCAL, **kwargs)
100
101 def __init__(self,
102 forward_to,
103 src_port,
104 dst_port,
105 user,
106 identity_file,
107 host,
108 port=22,
109 connect_timeout=_DEFAULT_CONNECT_TIMEOUT,
110 alive_interval=_DEFAULT_ALIVE_INTERVAL,
111 disconnect_wait=_DEFAULT_DISCONNECT_WAIT,
112 retries=_DEFAULT_RETRIES,
113 exp_factor=_DEFAULT_EXP_FACTOR):
114 """Constructor.
115
116 Args:
117 forward_to: Which direction to forward traffic: REMOTE or LOCAL.
118 src_port: Source port for forwarding.
119 dst_port: Destination port for forwarding.
120 user: Username on remote server.
121 identity_file: Identity file for passwordless authentication on remote
122 server.
123 host: Host of remote server.
124 port: Port of remote server.
125 connect_timeout: Time in seconds
126 alive_interval:
127 disconnect_wait: The number of seconds to wait before reconnecting after
128 the first disconnect.
129 retries: The number of times to retry before reporting a failed
130 connection.
131 exp_factor: After each reconnect, the disconnect wait time is multiplied
132 by 2^exp_factor.
133 """
134 # Internal use.
135 self._ssh_thread = None
136 self._ssh_output = None
137 self._exception = None
138 self._state = self.CONNECTING
139 self._poll = threading.Event()
140
141 # Connection arguments.
142 self._forward_to = forward_to
143 self._src_port = src_port
144 self._dst_port = dst_port
145 self._host = host
146 self._user = user
147 self._identity_file = identity_file
148 self._port = port
149
150 # Configuration arguments.
151 self._connect_timeout = connect_timeout
152 self._alive_interval = alive_interval
153 self._exp_factor = exp_factor
154
155 t = threading.Thread(
156 target=self._Run,
157 args=(disconnect_wait, retries))
158 t.daemon = True
159 t.start()
160
161 def __str__(self):
162 # State representation.
163 if self._state == self.CONNECTING:
164 state_str = 'connecting'
165 elif self._state == self.INITIALIZED:
166 state_str = 'initialized'
167 else:
168 state_str = 'failed'
169
170 # Port forward representation.
171 if self._forward_to == self.REMOTE:
172 fwd_str = '->%d' % self._dst_port
173 else:
174 fwd_str = '%d<-' % self._dst_port
175
176 return 'SSHPortForwarder(%s,%s)' % (state_str, fwd_str)
177
178 def _ForwardArgs(self):
179 if self._forward_to == self.REMOTE:
180 return ['-R', '%d:127.0.0.1:%d' % (self._dst_port, self._src_port)]
181 else:
182 return ['-L', '%d:127.0.0.1:%d' % (self._src_port, self._dst_port)]
183
184 def _RunSSHCmd(self):
185 """Runs the SSH command, storing the exception on failure."""
186 try:
187 cmd = [
188 'ssh',
189 '-o', 'StrictHostKeyChecking=no',
190 '-o', 'GlobalKnownHostsFile=/dev/null',
191 '-o', 'UserKnownHostsFile=/dev/null',
192 '-o', 'ExitOnForwardFailure=yes',
193 '-o', 'ConnectTimeout=%d' % self._connect_timeout,
194 '-o', 'ServerAliveInterval=%d' % self._alive_interval,
195 '-o', 'ServerAliveCountMax=1',
196 '-o', 'TCPKeepAlive=yes',
197 '-o', 'BatchMode=yes',
198 '-i', self._identity_file,
199 '-N',
200 '-p', str(self._port),
201 '%s@%s' % (self._user, self._host),
202 ] + self._ForwardArgs()
203 logging.info(' '.join(cmd))
204 self._ssh_output = subprocess.check_output(cmd, stderr=subprocess.STDOUT)
205 except subprocess.CalledProcessError as e:
206 self._exception = e
207 finally:
208 pass
209
210 def _Run(self, disconnect_wait, retries):
211 """Wraps around the SSH command, detecting its connection status."""
212 assert retries > 0, '%s: _Run must be called with retries > 0' % self
213
214 logging.info('%s: Connecting to %s:%d',
215 self, self._host, self._port)
216
217 # Set identity file permissions. Need to only be user-readable for ssh to
218 # use the key.
219 try:
220 os.chmod(self._identity_file, 0600)
221 except OSError as e:
222 logging.error('%s: Error setting identity file permissions: %s',
223 self, e)
224 self._state = self.FAILED
225 return
226
227 # Start a thread. If it fails, deal with the failure. If it is still
228 # running after connect_timeout seconds, assume everything's working great,
229 # and tell the caller. Then, continue waiting for it to end.
230 self._ssh_thread = threading.Thread(target=self._RunSSHCmd)
231 self._ssh_thread.daemon = True
232 self._ssh_thread.start()
233
234 # See if the SSH thread is still working after connect_timeout.
235 self._ssh_thread.join(self._connect_timeout)
236 if self._ssh_thread.is_alive():
237 # Assumed to be working. Tell our caller that we are connected.
238 if self._state != self.INITIALIZED:
239 self._state = self.INITIALIZED
240 self._poll.set()
241 logging.info('%s: Still connected after timeout=%ds',
242 self, self._connect_timeout)
243
244 # Only for debug purposes. Keep showing connection status.
245 while self._ssh_thread.is_alive():
246 logging.debug('%s: Still connected', self)
247 self._ssh_thread.join(self._DEBUG_INTERVAL)
248
249 # Figure out what went wrong.
250 if not self._exception:
251 logging.info('%s: SSH unexpectedly exited: %s',
252 self, self._ssh_output.rstrip())
253 if self._exception and self._FAILED_STR in self._exception.output:
254 self._state = self.FAILED
255 self._poll.set()
256 logging.info('%s: Port forwarding failed', self)
257 return
258 elif retries == 1:
259 self._state = self.FAILED
260 self._poll.set()
261 logging.info('%s: Disconnected (0 retries left)', self)
262 return
263 else:
264 logging.info('%s: Disconnected, retrying (sleep %1ds, %d retries left)',
265 self, disconnect_wait, retries - 1)
266 time.sleep(disconnect_wait)
267 self._Run(disconnect_wait=disconnect_wait * (2 ** self._exp_factor),
268 retries=retries - 1)
269
270 def GetState(self):
271 """Returns the current connection state.
272
273 State may be one of:
274
275 CONNECTING: Still attempting to make the first successful connection.
276 INITIALIZED: Is either connected or is trying to make subsequent
277 connection.
278 FAILED: Has completed all connection attempts, or server has reported that
279 target port is in use.
280 """
281 return self._state
282
283 def GetDstPort(self):
284 """Returns the current target port."""
285 return self._dst_port
286
287 def Wait(self):
288 """Waits for a state change, and returns the new state."""
289 self._poll.wait()
290 self._poll.clear()
291 return self.GetState()
292
293
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800294class Ghost(object):
295 """Ghost implements the client protocol of Overlord.
296
297 Ghost provide terminal/shell/logcat functionality and manages the client
298 side connectivity.
299 """
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800300 NONE, AGENT, TERMINAL, SHELL, LOGCAT, FILE = range(6)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800301
302 MODE_NAME = {
303 NONE: 'NONE',
304 AGENT: 'Agent',
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800305 TERMINAL: 'Terminal',
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800306 SHELL: 'Shell',
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800307 LOGCAT: 'Logcat',
308 FILE: 'File'
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800309 }
310
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +0800311 RANDOM_MID = '##random_mid##'
312
Wei-Ning Huangd521f282015-08-07 05:28:04 +0800313 def __init__(self, overlord_addrs, mode=AGENT, mid=None, sid=None,
Wei-Ning Huang23ed0162015-09-18 14:42:03 +0800314 prop_file=None, terminal_sid=None, tty_device=None,
315 command=None, file_op=None):
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800316 """Constructor.
317
318 Args:
319 overlord_addrs: a list of possible address of overlord.
320 mode: client mode, either AGENT, SHELL or LOGCAT
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +0800321 mid: a str to set for machine ID. If mid equals Ghost.RANDOM_MID, machine
322 id is randomly generated.
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800323 sid: session ID. If the connection is requested by overlord, sid should
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800324 be set to the corresponding session id assigned by overlord.
Wei-Ning Huang23ed0162015-09-18 14:42:03 +0800325 prop_file: properties file filename.
Wei-Ning Huangd521f282015-08-07 05:28:04 +0800326 terminal_sid: the terminal session ID associate with this client. This is
327 use for file download.
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800328 tty_device: the terminal device to open, if tty_device is None, as pseudo
329 terminal will be opened instead.
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800330 command: the command to execute when we are in SHELL mode.
Wei-Ning Huange2981862015-08-03 15:03:08 +0800331 file_op: a tuple (action, filepath, pid). action is either 'download' or
332 'upload'. pid is the pid of the target shell, used to determine where
333 the current working is and thus where to upload to.
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800334 """
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800335 assert mode in [Ghost.AGENT, Ghost.TERMINAL, Ghost.SHELL, Ghost.FILE]
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800336 if mode == Ghost.SHELL:
337 assert command is not None
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800338 if mode == Ghost.FILE:
339 assert file_op is not None
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800340
341 self._overlord_addrs = overlord_addrs
Wei-Ning Huangad330c52015-03-12 20:34:18 +0800342 self._connected_addr = None
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800343 self._mode = mode
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +0800344 self._mid = mid
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800345 self._sock = None
346 self._machine_id = self.GetMachineID()
Wei-Ning Huangfed95862015-08-07 03:17:11 +0800347 self._session_id = sid if sid is not None else str(uuid.uuid4())
Wei-Ning Huangd521f282015-08-07 05:28:04 +0800348 self._terminal_session_id = terminal_sid
Wei-Ning Huang23ed0162015-09-18 14:42:03 +0800349 self._ttyname_to_sid = {}
350 self._terminal_sid_to_pid = {}
351 self._prop_file = prop_file
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800352 self._properties = {}
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800353 self._tty_device = tty_device
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800354 self._shell_command = command
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800355 self._file_op = file_op
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800356 self._buf = ''
357 self._requests = {}
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800358 self._reset = threading.Event()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800359 self._last_ping = 0
360 self._queue = Queue.Queue()
Joel Kitching22b89042015-08-06 18:23:29 +0800361 self._forward_ssh = False
362 self._ssh_port_forwarder = None
363 self._target_identity_file = os.path.join(os.path.dirname(
364 os.path.abspath(os.path.realpath(__file__))), 'ghost_rsa')
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800365 self._download_queue = Queue.Queue()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800366
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800367 def SetIgnoreChild(self, status):
368 # Only ignore child for Agent since only it could spawn child Ghost.
369 if self._mode == Ghost.AGENT:
370 signal.signal(signal.SIGCHLD,
371 signal.SIG_IGN if status else signal.SIG_DFL)
372
373 def GetFileSha1(self, filename):
374 with open(filename, 'r') as f:
375 return hashlib.sha1(f.read()).hexdigest()
376
Wei-Ning Huang58833882015-09-16 16:52:37 +0800377 def UseSSL(self):
378 """Determine if SSL is enabled on the Overlord server."""
379 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
380 try:
381 sock.connect((self._connected_addr[0], _OVERLORD_HTTP_PORT))
382 sock.send('GET\r\n')
383
384 data = sock.recv(16)
385 return 'HTTP' not in data
386 except Exception:
387 return False # For whatever reason above failed, assume HTTP
388
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800389 def Upgrade(self):
390 logging.info('Upgrade: initiating upgrade sequence...')
391
392 scriptpath = os.path.abspath(sys.argv[0])
Wei-Ning Huang03f9f762015-09-16 21:51:35 +0800393 url = 'http%s://%s:%d/upgrade/ghost.py' % (
394 's' if self.UseSSL() else '', self._connected_addr[0],
395 _OVERLORD_HTTP_PORT)
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800396
397 # Download sha1sum for ghost.py for verification
398 try:
399 with contextlib.closing(urllib.urlopen(url + '.sha1')) as f:
400 if f.getcode() != 200:
401 raise RuntimeError('HTTP status %d' % f.getcode())
402 sha1sum = f.read().strip()
403 except Exception:
404 logging.error('Upgrade: failed to download sha1sum file, abort')
405 return
406
407 if self.GetFileSha1(scriptpath) == sha1sum:
408 logging.info('Upgrade: ghost is already up-to-date, skipping upgrade')
409 return
410
411 # Download upgrade version of ghost.py
412 try:
413 with contextlib.closing(urllib.urlopen(url)) as f:
414 if f.getcode() != 200:
415 raise RuntimeError('HTTP status %d' % f.getcode())
416 data = f.read()
417 except Exception:
418 logging.error('Upgrade: failed to download upgrade, abort')
419 return
420
421 # Compare SHA1 sum
422 if hashlib.sha1(data).hexdigest() != sha1sum:
423 logging.error('Upgrade: sha1sum mismatch, abort')
424 return
425
426 python = os.readlink('/proc/self/exe')
427 try:
428 with open(scriptpath, 'w') as f:
429 f.write(data)
430 except Exception:
431 logging.error('Upgrade: failed to write upgrade onto disk, abort')
432 return
433
434 logging.info('Upgrade: restarting ghost...')
435 self.CloseSockets()
436 self.SetIgnoreChild(False)
437 os.execve(python, [python, scriptpath] + sys.argv[1:], os.environ)
438
Wei-Ning Huang23ed0162015-09-18 14:42:03 +0800439 def LoadProperties(self):
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800440 try:
Wei-Ning Huang23ed0162015-09-18 14:42:03 +0800441 if self._prop_file:
442 with open(self._prop_file, 'r') as f:
443 self._properties = json.loads(f.read())
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800444 except Exception as e:
Wei-Ning Huang23ed0162015-09-18 14:42:03 +0800445 logging.exception('LoadProperties: ' + str(e))
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800446
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800447 def CloseSockets(self):
448 # Close sockets opened by parent process, since we don't use it anymore.
449 for fd in os.listdir('/proc/self/fd/'):
450 try:
451 real_fd = os.readlink('/proc/self/fd/%s' % fd)
452 if real_fd.startswith('socket'):
453 os.close(int(fd))
454 except Exception:
455 pass
456
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800457 def SpawnGhost(self, mode, sid=None, terminal_sid=None, tty_device=None,
458 command=None, file_op=None):
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800459 """Spawn a child ghost with specific mode.
460
461 Returns:
462 The spawned child process pid.
463 """
Joel Kitching22b89042015-08-06 18:23:29 +0800464 # Restore the default signal handler, so our child won't have problems.
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800465 self.SetIgnoreChild(False)
466
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800467 pid = os.fork()
468 if pid == 0:
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800469 self.CloseSockets()
Wei-Ning Huangd521f282015-08-07 05:28:04 +0800470 g = Ghost([self._connected_addr], mode, Ghost.RANDOM_MID, sid,
Wei-Ning Huang23ed0162015-09-18 14:42:03 +0800471 terminal_sid=terminal_sid, tty_device=tty_device,
472 command=command, file_op=file_op)
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800473 g.Start()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800474 sys.exit(0)
475 else:
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800476 self.SetIgnoreChild(True)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800477 return pid
478
479 def Timestamp(self):
480 return int(time.time())
481
482 def GetGateWayIP(self):
483 with open('/proc/net/route', 'r') as f:
484 lines = f.readlines()
485
486 ips = []
487 for line in lines:
488 parts = line.split('\t')
489 if parts[2] == '00000000':
490 continue
491
492 try:
493 h = parts[2].decode('hex')
494 ips.append('%d.%d.%d.%d' % tuple(ord(x) for x in reversed(h)))
495 except TypeError:
496 pass
497
498 return ips
499
Wei-Ning Huang829e0c82015-05-26 14:37:23 +0800500 def GetShopfloorIP(self):
501 try:
502 import factory_common # pylint: disable=W0612
503 from cros.factory.test import shopfloor
504
505 url = shopfloor.get_server_url()
506 match = re.match(r'^https?://(.*):.*$', url)
507 if match:
508 return [match.group(1)]
509 except Exception:
510 pass
511 return []
512
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800513 def GetMachineID(self):
514 """Generates machine-dependent ID string for a machine.
515 There are many ways to generate a machine ID:
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800516 1. factory device_id
517 2. factory device-data
Wei-Ning Huang1d7603b2015-07-03 17:38:56 +0800518 3. /sys/class/dmi/id/product_uuid (only available on intel machines)
519 4. MAC address
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800520 We follow the listed order to generate machine ID, and fallback to the next
521 alternative if the previous doesn't work.
522 """
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +0800523 if self._mid == Ghost.RANDOM_MID:
Wei-Ning Huangaed90452015-03-23 17:50:21 +0800524 return str(uuid.uuid4())
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +0800525 elif self._mid:
526 return self._mid
Wei-Ning Huangaed90452015-03-23 17:50:21 +0800527
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800528 # Try factory device id
529 try:
530 import factory_common # pylint: disable=W0612
531 from cros.factory.test import event_log
532 with open(event_log.DEVICE_ID_PATH) as f:
533 return f.read().strip()
534 except Exception:
535 pass
536
Wei-Ning Huang1d7603b2015-07-03 17:38:56 +0800537 # Try factory device data
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800538 try:
539 p = subprocess.Popen('factory device-data | grep mlb_serial_number | '
540 'cut -d " " -f 2', stdout=subprocess.PIPE,
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800541 stderr=subprocess.PIPE, shell=True)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800542 stdout, _ = p.communicate()
543 if stdout == '':
Wei-Ning Huang1d7603b2015-07-03 17:38:56 +0800544 raise RuntimeError('empty mlb number')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800545 return stdout.strip()
546 except Exception:
547 pass
548
Wei-Ning Huang1d7603b2015-07-03 17:38:56 +0800549 # Try DMI product UUID
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800550 try:
551 with open('/sys/class/dmi/id/product_uuid', 'r') as f:
552 return f.read().strip()
553 except Exception:
554 pass
555
Wei-Ning Huang1d7603b2015-07-03 17:38:56 +0800556 # Use MAC address if non is available
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800557 try:
558 macs = []
559 ifaces = sorted(os.listdir('/sys/class/net'))
560 for iface in ifaces:
561 if iface == 'lo':
562 continue
563
564 with open('/sys/class/net/%s/address' % iface, 'r') as f:
565 macs.append(f.read().strip())
566
567 return ';'.join(macs)
568 except Exception:
569 pass
570
Moja Hsuc9ecc8b2015-07-13 11:39:17 +0800571 raise RuntimeError('can\'t generate machine ID')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800572
573 def Reset(self):
574 """Reset state and clear request handlers."""
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800575 self._reset.clear()
Moja Hsuc9ecc8b2015-07-13 11:39:17 +0800576 self._buf = ''
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800577 self._last_ping = 0
578 self._requests = {}
Wei-Ning Huang23ed0162015-09-18 14:42:03 +0800579 self.LoadProperties()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800580
581 def SendMessage(self, msg):
582 """Serialize the message and send it through the socket."""
583 self._sock.send(json.dumps(msg) + _SEPARATOR)
584
585 def SendRequest(self, name, args, handler=None,
586 timeout=_REQUEST_TIMEOUT_SECS):
587 if handler and not callable(handler):
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800588 raise RequestError('Invalid request handler for msg "%s"' % name)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800589
590 rid = str(uuid.uuid4())
591 msg = {'rid': rid, 'timeout': timeout, 'name': name, 'params': args}
Wei-Ning Huange2981862015-08-03 15:03:08 +0800592 if timeout >= 0:
593 self._requests[rid] = [self.Timestamp(), timeout, handler]
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800594 self.SendMessage(msg)
595
596 def SendResponse(self, omsg, status, params=None):
597 msg = {'rid': omsg['rid'], 'response': status, 'params': params}
598 self.SendMessage(msg)
599
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800600 def HandleTTYControl(self, fd, control_string):
Moja Hsuc9ecc8b2015-07-13 11:39:17 +0800601 msg = json.loads(control_string)
602 command = msg['command']
603 params = msg['params']
604 if command == 'resize':
605 # some error happened on websocket
606 if len(params) != 2:
607 return
608 winsize = struct.pack('HHHH', params[0], params[1], 0, 0)
609 fcntl.ioctl(fd, termios.TIOCSWINSZ, winsize)
610 else:
611 logging.warn('Invalid request command "%s"', command)
612
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800613 def SpawnTTYServer(self, _):
614 """Spawn a TTY server and forward I/O to the TCP socket."""
615 logging.info('SpawnTTYServer: started')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800616
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800617 try:
618 if self._tty_device is None:
619 pid, fd = os.forkpty()
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800620
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800621 if pid == 0:
622 ttyname = os.readlink('/proc/%d/fd/0' % os.getpid())
623 try:
624 server = GhostRPCServer()
625 server.RegisterTTY(self._session_id, ttyname)
626 server.RegisterSession(self._session_id, os.getpid())
627 except Exception:
628 # If ghost is launched without RPC server, the call will fail but we
629 # can ignore it.
630 pass
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800631
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800632 # The directory that contains the current running ghost script
633 script_dir = os.path.dirname(os.path.abspath(sys.argv[0]))
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800634
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800635 env = os.environ.copy()
636 env['USER'] = os.getenv('USER', 'root')
637 env['HOME'] = os.getenv('HOME', '/root')
638 env['PATH'] = os.getenv('PATH') + ':%s' % script_dir
639 os.chdir(env['HOME'])
640 os.execve(_SHELL, [_SHELL], env)
641 else:
642 fd = os.open(self._tty_device, os.O_RDWR)
Wei-Ning Huang39169902015-09-19 06:00:23 +0800643 tty.setraw(fd)
644 attr = termios.tcgetattr(fd)
645 attr[0] &= ~(termios.IXON | termios.IXOFF)
646 attr[2] |= termios.CLOCAL
647 attr[2] &= ~termios.CRTSCTS
648 attr[4] = termios.B115200
649 attr[5] = termios.B115200
650 termios.tcsetattr(fd, termios.TCSANOW, attr)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800651
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800652 control_state = None
653 control_string = ''
654 write_buffer = ''
655 while True:
656 rd, _, _ = select.select([self._sock, fd], [], [])
657
658 if fd in rd:
659 self._sock.send(os.read(fd, _BUFSIZE))
660
661 if self._sock in rd:
662 ret = self._sock.recv(_BUFSIZE)
663 if len(ret) == 0:
664 raise RuntimeError('socket closed')
665 while ret:
666 if control_state:
667 if chr(_CONTROL_END) in ret:
668 index = ret.index(chr(_CONTROL_END))
669 control_string += ret[:index]
670 self.HandleTTYControl(fd, control_string)
671 control_state = None
672 control_string = ''
673 ret = ret[index+1:]
Moja Hsuc9ecc8b2015-07-13 11:39:17 +0800674 else:
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800675 control_string += ret
676 ret = ''
677 else:
678 if chr(_CONTROL_START) in ret:
679 control_state = _CONTROL_START
680 index = ret.index(chr(_CONTROL_START))
681 write_buffer += ret[:index]
682 ret = ret[index+1:]
683 else:
684 write_buffer += ret
685 ret = ''
686 if write_buffer:
687 os.write(fd, write_buffer)
688 write_buffer = ''
689 except (OSError, socket.error, RuntimeError) as e:
690 self._sock.close()
691 logging.info('SpawnTTYServer: %s' % str(e))
692 logging.info('SpawnTTYServer: terminated')
693 sys.exit(0)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800694
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800695 def SpawnShellServer(self, _):
696 """Spawn a shell server and forward input/output from/to the TCP socket."""
697 logging.info('SpawnShellServer: started')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800698
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800699 p = subprocess.Popen(self._shell_command, stdin=subprocess.PIPE,
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800700 stdout=subprocess.PIPE, stderr=subprocess.PIPE,
701 shell=True)
702
703 def make_non_block(fd):
704 fl = fcntl.fcntl(fd, fcntl.F_GETFL)
705 fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
706
707 make_non_block(p.stdout)
708 make_non_block(p.stderr)
709
710 try:
711 while True:
712 rd, _, _ = select.select([p.stdout, p.stderr, self._sock], [], [])
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800713 if p.stdout in rd:
714 self._sock.send(p.stdout.read(_BUFSIZE))
715
716 if p.stderr in rd:
717 self._sock.send(p.stderr.read(_BUFSIZE))
718
719 if self._sock in rd:
720 ret = self._sock.recv(_BUFSIZE)
721 if len(ret) == 0:
Moja Hsuc9ecc8b2015-07-13 11:39:17 +0800722 raise RuntimeError('socket closed')
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800723 p.stdin.write(ret)
Wei-Ning Huangf14c84e2015-08-03 15:03:08 +0800724 p.poll()
725 if p.returncode != None:
726 break
727 finally:
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800728 self._sock.close()
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800729 logging.info('SpawnShellServer: terminated')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800730 sys.exit(0)
731
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800732 def InitiateFileOperation(self, _):
733 if self._file_op[0] == 'download':
734 size = os.stat(self._file_op[1]).st_size
735 self.SendRequest('request_to_download',
Wei-Ning Huangd521f282015-08-07 05:28:04 +0800736 {'terminal_sid': self._terminal_session_id,
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800737 'filename': os.path.basename(self._file_op[1]),
738 'size': size})
Wei-Ning Huange2981862015-08-03 15:03:08 +0800739 elif self._file_op[0] == 'upload':
740 self.SendRequest('clear_to_upload', {}, timeout=-1)
741 self.StartUploadServer()
742 else:
743 logging.error('InitiateFileOperation: unknown file operation, ignored')
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800744
745 def StartDownloadServer(self):
746 logging.info('StartDownloadServer: started')
747
748 try:
749 with open(self._file_op[1], 'rb') as f:
750 while True:
751 data = f.read(_BLOCK_SIZE)
752 if len(data) == 0:
753 break
754 self._sock.send(data)
755 except Exception as e:
756 logging.error('StartDownloadServer: %s', e)
757 finally:
758 self._sock.close()
759
760 logging.info('StartDownloadServer: terminated')
761 sys.exit(0)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800762
Wei-Ning Huange2981862015-08-03 15:03:08 +0800763 def StartUploadServer(self):
764 logging.info('StartUploadServer: started')
765
766 try:
767 target_dir = os.getenv('HOME', '/tmp')
768
769 # Get the client's working dir, which is our target upload dir
770 if self._file_op[2]:
771 target_dir = os.readlink('/proc/%d/cwd' % self._file_op[2])
772
773 self._sock.setblocking(False)
774 with open(os.path.join(target_dir, self._file_op[1]), 'wb') as f:
775 while True:
776 rd, _, _ = select.select([self._sock], [], [])
777 if self._sock in rd:
778 buf = self._sock.recv(_BLOCK_SIZE)
779 if len(buf) == 0:
780 break
781 f.write(buf)
782 except socket.error as e:
783 logging.error('StartUploadServer: socket error: %s', e)
784 except Exception as e:
785 logging.error('StartUploadServer: %s', e)
786 finally:
787 self._sock.close()
788
789 logging.info('StartUploadServer: terminated')
790 sys.exit(0)
791
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800792 def Ping(self):
793 def timeout_handler(x):
794 if x is None:
795 raise PingTimeoutError
796
797 self._last_ping = self.Timestamp()
798 self.SendRequest('ping', {}, timeout_handler, 5)
799
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800800 def HandleRequest(self, msg):
Wei-Ning Huange2981862015-08-03 15:03:08 +0800801 command = msg['name']
802 params = msg['params']
803
804 if command == 'upgrade':
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800805 self.Upgrade()
Wei-Ning Huange2981862015-08-03 15:03:08 +0800806 elif command == 'terminal':
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800807 self.SpawnGhost(self.TERMINAL, params['sid'],
808 tty_device=params['tty_device'])
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800809 self.SendResponse(msg, RESPONSE_SUCCESS)
Wei-Ning Huange2981862015-08-03 15:03:08 +0800810 elif command == 'shell':
811 self.SpawnGhost(self.SHELL, params['sid'], command=params['command'])
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800812 self.SendResponse(msg, RESPONSE_SUCCESS)
Wei-Ning Huange2981862015-08-03 15:03:08 +0800813 elif command == 'file_download':
814 self.SpawnGhost(self.FILE, params['sid'],
815 file_op=('download', params['filename'], None))
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800816 self.SendResponse(msg, RESPONSE_SUCCESS)
Wei-Ning Huange2981862015-08-03 15:03:08 +0800817 elif command == 'clear_to_download':
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800818 self.StartDownloadServer()
Wei-Ning Huange2981862015-08-03 15:03:08 +0800819 elif command == 'file_upload':
820 pid = self._terminal_sid_to_pid.get(params['terminal_sid'], None)
821 self.SpawnGhost(self.FILE, params['sid'],
822 file_op=('upload', params['filename'], pid))
823 self.SendResponse(msg, RESPONSE_SUCCESS)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800824
825 def HandleResponse(self, response):
826 rid = str(response['rid'])
827 if rid in self._requests:
828 handler = self._requests[rid][2]
829 del self._requests[rid]
830 if callable(handler):
831 handler(response)
832 else:
833 print(response, self._requests.keys())
Joel Kitching22b89042015-08-06 18:23:29 +0800834 logging.warning('Received unsolicited response, ignored')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800835
836 def ParseMessage(self):
837 msgs_json = self._buf.split(_SEPARATOR)
838 self._buf = msgs_json.pop()
839
840 for msg_json in msgs_json:
841 try:
842 msg = json.loads(msg_json)
843 except ValueError:
844 # Ignore mal-formed message.
845 continue
846
847 if 'name' in msg:
848 self.HandleRequest(msg)
849 elif 'response' in msg:
850 self.HandleResponse(msg)
851 else: # Ingnore mal-formed message.
852 pass
853
854 def ScanForTimeoutRequests(self):
Joel Kitching22b89042015-08-06 18:23:29 +0800855 """Scans for pending requests which have timed out.
856
857 If any timed-out requests are discovered, their handler is called with the
858 special response value of None.
859 """
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800860 for rid in self._requests.keys()[:]:
861 request_time, timeout, handler = self._requests[rid]
862 if self.Timestamp() - request_time > timeout:
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800863 if callable(handler):
864 handler(None)
865 else:
866 logging.error('Request %s timeout', rid)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800867 del self._requests[rid]
868
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800869 def InitiateDownload(self):
870 ttyname, filename = self._download_queue.get()
Wei-Ning Huangd521f282015-08-07 05:28:04 +0800871 sid = self._ttyname_to_sid[ttyname]
872 self.SpawnGhost(self.FILE, terminal_sid=sid,
873 file_op=('download', filename, None))
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800874
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800875 def Listen(self):
876 try:
877 while True:
878 rds, _, _ = select.select([self._sock], [], [], _PING_INTERVAL / 2)
879
880 if self._sock in rds:
881 self._buf += self._sock.recv(_BUFSIZE)
882 self.ParseMessage()
883
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800884 if (self._mode == self.AGENT and
885 self.Timestamp() - self._last_ping > _PING_INTERVAL):
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800886 self.Ping()
887 self.ScanForTimeoutRequests()
888
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800889 if not self._download_queue.empty():
890 self.InitiateDownload()
891
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800892 if self._reset.is_set():
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800893 self.Reset()
894 break
895 except socket.error:
896 raise RuntimeError('Connection dropped')
897 except PingTimeoutError:
898 raise RuntimeError('Connection timeout')
899 finally:
900 self._sock.close()
901
902 self._queue.put('resume')
903
904 if self._mode != Ghost.AGENT:
905 sys.exit(1)
906
907 def Register(self):
908 non_local = {}
909 for addr in self._overlord_addrs:
910 non_local['addr'] = addr
911 def registered(response):
912 if response is None:
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800913 self._reset.set()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800914 raise RuntimeError('Register request timeout')
Wei-Ning Huang63c16092015-09-18 16:20:27 +0800915
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800916 logging.info('Registered with Overlord at %s:%d', *non_local['addr'])
Wei-Ning Huang63c16092015-09-18 16:20:27 +0800917 self._connected_addr = addr
918 self.Upgrade() # Check for upgrade
919 self._queue.put('pause', True)
920
Joel Kitching22b89042015-08-06 18:23:29 +0800921 if self._forward_ssh:
922 logging.info('Starting target SSH port negotiation')
923 self.NegotiateTargetSSHPort()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800924
925 try:
926 logging.info('Trying %s:%d ...', *addr)
927 self.Reset()
928 self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
929 self._sock.settimeout(_PING_TIMEOUT)
930 self._sock.connect(addr)
931
932 logging.info('Connection established, registering...')
933 handler = {
934 Ghost.AGENT: registered,
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800935 Ghost.TERMINAL: self.SpawnTTYServer,
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800936 Ghost.SHELL: self.SpawnShellServer,
937 Ghost.FILE: self.InitiateFileOperation,
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800938 }[self._mode]
939
940 # Machine ID may change if MAC address is used (USB-ethernet dongle
941 # plugged/unplugged)
942 self._machine_id = self.GetMachineID()
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800943 self.SendRequest('register',
944 {'mode': self._mode, 'mid': self._machine_id,
Wei-Ning Huangfed95862015-08-07 03:17:11 +0800945 'sid': self._session_id,
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800946 'properties': self._properties}, handler)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800947 except socket.error:
948 pass
949 else:
950 self._sock.settimeout(None)
951 self.Listen()
952
Moja Hsuc9ecc8b2015-07-13 11:39:17 +0800953 raise RuntimeError('Cannot connect to any server')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800954
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800955 def Reconnect(self):
956 logging.info('Received reconnect request from RPC server, reconnecting...')
957 self._reset.set()
958
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800959 def AddToDownloadQueue(self, ttyname, filename):
960 self._download_queue.put((ttyname, filename))
961
Wei-Ning Huangd521f282015-08-07 05:28:04 +0800962 def RegisterTTY(self, session_id, ttyname):
963 self._ttyname_to_sid[ttyname] = session_id
Wei-Ning Huange2981862015-08-03 15:03:08 +0800964
965 def RegisterSession(self, session_id, process_id):
966 self._terminal_sid_to_pid[session_id] = process_id
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800967
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800968 def StartLanDiscovery(self):
969 """Start to listen to LAN discovery packet at
970 _OVERLORD_LAN_DISCOVERY_PORT."""
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800971
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800972 def thread_func():
973 s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
974 s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
975 s.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800976 try:
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800977 s.bind(('0.0.0.0', _OVERLORD_LAN_DISCOVERY_PORT))
978 except socket.error as e:
Moja Hsuc9ecc8b2015-07-13 11:39:17 +0800979 logging.error('LAN discovery: %s, abort', e)
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800980 return
981
982 logging.info('LAN Discovery: started')
983 while True:
984 rd, _, _ = select.select([s], [], [], 1)
985
986 if s in rd:
987 data, source_addr = s.recvfrom(_BUFSIZE)
988 parts = data.split()
989 if parts[0] == 'OVERLORD':
990 ip, port = parts[1].split(':')
991 if not ip:
992 ip = source_addr[0]
993 self._queue.put((ip, int(port)), True)
994
995 try:
996 obj = self._queue.get(False)
997 except Queue.Empty:
998 pass
999 else:
1000 if type(obj) is not str:
1001 self._queue.put(obj)
1002 elif obj == 'pause':
1003 logging.info('LAN Discovery: paused')
1004 while obj != 'resume':
1005 obj = self._queue.get(True)
1006 logging.info('LAN Discovery: resumed')
1007
1008 t = threading.Thread(target=thread_func)
1009 t.daemon = True
1010 t.start()
1011
1012 def StartRPCServer(self):
Joel Kitching22b89042015-08-06 18:23:29 +08001013 logging.info('RPC Server: started')
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001014 rpc_server = SimpleJSONRPCServer((_DEFAULT_BIND_ADDRESS, _GHOST_RPC_PORT),
1015 logRequests=False)
1016 rpc_server.register_function(self.Reconnect, 'Reconnect')
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001017 rpc_server.register_function(self.RegisterTTY, 'RegisterTTY')
Wei-Ning Huange2981862015-08-03 15:03:08 +08001018 rpc_server.register_function(self.RegisterSession, 'RegisterSession')
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001019 rpc_server.register_function(self.AddToDownloadQueue, 'AddToDownloadQueue')
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001020 t = threading.Thread(target=rpc_server.serve_forever)
1021 t.daemon = True
1022 t.start()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001023
Wei-Ning Huang829e0c82015-05-26 14:37:23 +08001024 def ScanServer(self):
1025 for meth in [self.GetGateWayIP, self.GetShopfloorIP]:
1026 for addr in [(x, _OVERLORD_PORT) for x in meth()]:
1027 if addr not in self._overlord_addrs:
1028 self._overlord_addrs.append(addr)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001029
Joel Kitching22b89042015-08-06 18:23:29 +08001030 def NegotiateTargetSSHPort(self):
1031 """Request-receive target SSH port forwarding loop.
1032
1033 Repeatedly attempts to forward this machine's SSH port to target. It
1034 bounces back and forth between RequestPort and ReceivePort when a new port
1035 is required. ReceivePort starts a new thread so that the main ghost thread
1036 may continue running.
1037 """
1038 # Sanity check for identity file.
1039 if not os.path.isfile(self._target_identity_file):
1040 logging.info('No target host identity file: not negotiating '
1041 'target SSH port')
1042 return
1043
1044 def PollSSHPortForwarder():
1045 def ThreadFunc():
1046 while True:
1047 state = self._ssh_port_forwarder.GetState()
1048
1049 # Connected successfully.
1050 if state == SSHPortForwarder.INITIALIZED:
1051 # The SSH port forward has succeeded! Let's tell Overlord.
1052 port = self._ssh_port_forwarder.GetDstPort()
1053 RegisterPort(port)
1054
1055 # We've given up... continue to the next port.
1056 elif state == SSHPortForwarder.FAILED:
1057 break
1058
1059 # Either CONNECTING or INITIALIZED.
1060 self._ssh_port_forwarder.Wait()
1061
1062 # Only request a new port if we are still registered to Overlord.
1063 # Otherwise, a new call to NegotiateTargetSSHPort will be made,
1064 # which will take care of it.
1065 try:
1066 RequestPort()
1067 except Exception:
1068 logging.info('Failed to request port, will wait for next connection')
1069 self._ssh_port_forwarder = None
1070
1071 t = threading.Thread(target=ThreadFunc)
1072 t.daemon = True
1073 t.start()
1074
1075 def ReceivePort(response):
1076 # If the response times out, this version of Overlord may not support SSH
1077 # port negotiation. Give up on port negotiation process.
1078 if response is None:
1079 return
1080
1081 port = int(response['params']['port'])
1082 logging.info('Received target SSH port: %d', port)
1083
1084 if (self._ssh_port_forwarder and
1085 self._ssh_port_forwarder.GetState() != SSHPortForwarder.FAILED):
1086 logging.info('Unexpectedly received a target SSH port')
1087 return
1088
1089 # Try forwarding SSH port to target.
1090 self._ssh_port_forwarder = SSHPortForwarder.ToRemote(
1091 src_port=22,
1092 dst_port=port,
1093 user='ghost',
1094 identity_file=self._target_identity_file,
1095 host=self._connected_addr[0]) # Use Overlord host as target.
1096
1097 # Creates a new thread.
1098 PollSSHPortForwarder()
1099
1100 def RequestPort():
1101 logging.info('Requesting new target SSH port')
1102 self.SendRequest('request_target_ssh_port', {}, ReceivePort, 5)
1103
1104 def RegisterPort(port):
1105 logging.info('Registering target SSH port %d', port)
1106 self.SendRequest(
1107 'register_target_ssh_port',
1108 {'port': port}, RegisterPortResponse, 5)
1109
1110 def RegisterPortResponse(response):
1111 # Overlord responded to request_port already. If register_port fails,
1112 # something might be in an inconsistent state, so trigger a reconnect
1113 # via PingTimeoutError.
1114 if response is None:
1115 raise PingTimeoutError
1116 logging.info('Registering target SSH port acknowledged')
1117
1118 # If the SSHPortForwarder is already in a INITIALIZED state, we need to
1119 # manually report the port to target, since SSHPortForwarder is currently
1120 # blocking.
1121 if (self._ssh_port_forwarder and
1122 self._ssh_port_forwarder.GetState() == SSHPortForwarder.INITIALIZED):
1123 RegisterPort(self._ssh_port_forwarder.GetDstPort())
1124 if not self._ssh_port_forwarder:
1125 RequestPort()
1126
1127 def Start(self, lan_disc=False, rpc_server=False, forward_ssh=False):
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001128 logging.info('%s started', self.MODE_NAME[self._mode])
1129 logging.info('MID: %s', self._machine_id)
Wei-Ning Huangfed95862015-08-07 03:17:11 +08001130 logging.info('SID: %s', self._session_id)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001131
Wei-Ning Huangb05cde32015-08-01 09:48:41 +08001132 # We don't care about child process's return code, not wait is needed. This
1133 # is used to prevent zombie process from lingering in the system.
1134 self.SetIgnoreChild(True)
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001135
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001136 if lan_disc:
1137 self.StartLanDiscovery()
1138
1139 if rpc_server:
1140 self.StartRPCServer()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001141
Joel Kitching22b89042015-08-06 18:23:29 +08001142 self._forward_ssh = forward_ssh
1143
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001144 try:
1145 while True:
1146 try:
1147 addr = self._queue.get(False)
1148 except Queue.Empty:
1149 pass
1150 else:
1151 if type(addr) == tuple and addr not in self._overlord_addrs:
1152 logging.info('LAN Discovery: got overlord address %s:%d', *addr)
1153 self._overlord_addrs.append(addr)
1154
1155 try:
Wei-Ning Huang829e0c82015-05-26 14:37:23 +08001156 self.ScanServer()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001157 self.Register()
Joel Kitching22b89042015-08-06 18:23:29 +08001158 # Don't show stack trace for RuntimeError, which we use in this file for
1159 # plausible and expected errors (such as can't connect to server).
1160 except RuntimeError as e:
1161 logging.info('%s: %s, retrying in %ds',
1162 e.__class__.__name__, e.message, _RETRY_INTERVAL)
1163 time.sleep(_RETRY_INTERVAL)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001164 except Exception as e:
Joel Kitching22b89042015-08-06 18:23:29 +08001165 _, _, exc_traceback = sys.exc_info()
1166 traceback.print_tb(exc_traceback)
1167 logging.info('%s: %s, retrying in %ds',
1168 e.__class__.__name__, e.message, _RETRY_INTERVAL)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001169 time.sleep(_RETRY_INTERVAL)
1170
1171 self.Reset()
1172 except KeyboardInterrupt:
1173 logging.error('Received keyboard interrupt, quit')
1174 sys.exit(0)
1175
1176
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001177def GhostRPCServer():
Wei-Ning Huang8037c182015-09-19 04:41:50 +08001178 """Returns handler to Ghost's JSON RPC server."""
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001179 return jsonrpclib.Server('http://localhost:%d' % _GHOST_RPC_PORT)
1180
1181
Wei-Ning Huang8037c182015-09-19 04:41:50 +08001182def ForkToBackground():
1183 """Fork process to run in background."""
1184 pid = os.fork()
1185 if pid != 0:
1186 logging.info('Ghost(%d) running in background.', pid)
1187 sys.exit(0)
1188
1189
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001190def DownloadFile(filename):
Wei-Ning Huang8037c182015-09-19 04:41:50 +08001191 """Initiate a client-initiated file download."""
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001192 filepath = os.path.abspath(filename)
1193 if not os.path.exists(filepath):
Joel Kitching22b89042015-08-06 18:23:29 +08001194 logging.error('file `%s\' does not exist', filename)
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001195 sys.exit(1)
1196
1197 # Check if we actually have permission to read the file
1198 if not os.access(filepath, os.R_OK):
Joel Kitching22b89042015-08-06 18:23:29 +08001199 logging.error('can not open %s for reading', filepath)
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001200 sys.exit(1)
1201
1202 server = GhostRPCServer()
1203 server.AddToDownloadQueue(os.ttyname(0), filepath)
1204 sys.exit(0)
1205
1206
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001207def main():
1208 logger = logging.getLogger()
1209 logger.setLevel(logging.INFO)
1210
Wei-Ning Huang7d029b12015-03-06 10:32:15 +08001211 parser = argparse.ArgumentParser()
Wei-Ning Huang8037c182015-09-19 04:41:50 +08001212 parser.add_argument('--fork', dest='fork', action='store_true', default=False,
1213 help='fork procecess to run in background')
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +08001214 parser.add_argument('--mid', metavar='MID', dest='mid', action='store',
1215 default=None, help='use MID as machine ID')
1216 parser.add_argument('--rand-mid', dest='mid', action='store_const',
1217 const=Ghost.RANDOM_MID, help='use random machine ID')
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001218 parser.add_argument('--no-lan-disc', dest='lan_disc', action='store_false',
1219 default=True, help='disable LAN discovery')
1220 parser.add_argument('--no-rpc-server', dest='rpc_server',
1221 action='store_false', default=True,
1222 help='disable RPC server')
Wei-Ning Huang03f9f762015-09-16 21:51:35 +08001223 parser.add_argument('--forward-ssh', dest='forward_ssh',
1224 action='store_true', default=False,
1225 help='enable target SSH port forwarding')
Joel Kitching22b89042015-08-06 18:23:29 +08001226 parser.add_argument('--prop-file', metavar='PROP_FILE', dest='prop_file',
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001227 type=str, default=None,
Wei-Ning Huang7d029b12015-03-06 10:32:15 +08001228 help='file containing the JSON representation of client '
1229 'properties')
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001230 parser.add_argument('--download', metavar='FILE', dest='download', type=str,
1231 default=None, help='file to download')
Wei-Ning Huang23ed0162015-09-18 14:42:03 +08001232 parser.add_argument('--reset', dest='reset', default=False,
1233 action='store_true',
1234 help='reset ghost and reload all configs')
Wei-Ning Huang7d029b12015-03-06 10:32:15 +08001235 parser.add_argument('overlord_ip', metavar='OVERLORD_IP', type=str,
1236 nargs='*', help='overlord server address')
1237 args = parser.parse_args()
1238
Wei-Ning Huang8037c182015-09-19 04:41:50 +08001239 if args.fork:
1240 ForkToBackground()
1241
Wei-Ning Huang23ed0162015-09-18 14:42:03 +08001242 if args.reset:
1243 GhostRPCServer().Reconnect()
1244 sys.exit()
1245
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001246 if args.download:
1247 DownloadFile(args.download)
1248
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001249 addrs = [('localhost', _OVERLORD_PORT)]
Wei-Ning Huang7d029b12015-03-06 10:32:15 +08001250 addrs += [(x, _OVERLORD_PORT) for x in args.overlord_ip]
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001251
Wei-Ning Huang23ed0162015-09-18 14:42:03 +08001252 g = Ghost(addrs, Ghost.AGENT, args.mid, prop_file=args.prop_file)
Joel Kitching22b89042015-08-06 18:23:29 +08001253 g.Start(args.lan_disc, args.rpc_server, args.forward_ssh)
1254
1255
1256def _SigtermHandler(*_):
1257 """Ensure that SSH processes also get killed on a sigterm signal.
1258
1259 By also passing the sigterm signal onto the process group, we ensure that any
1260 child SSH processes will also get killed.
1261
1262 Source:
1263 http://www.tsheffler.com/blog/2010/11/21/python-multithreaded-daemon-with-sigterm-support-a-recipe/
1264 """
1265 logging.info('SIGTERM handler: shutting down')
1266 if not _SigtermHandler.SIGTERM_SENT:
1267 _SigtermHandler.SIGTERM_SENT = True
1268 logging.info('Sending TERM to process group')
1269 os.killpg(0, signal.SIGTERM)
1270 sys.exit()
1271_SigtermHandler.SIGTERM_SENT = False
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001272
1273
1274if __name__ == '__main__':
Joel Kitching22b89042015-08-06 18:23:29 +08001275 signal.signal(signal.SIGTERM, _SigtermHandler)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001276 main()