blob: a476375d876f05147b66f6cf09340620535ea3f6 [file] [log] [blame]
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +08001#!/usr/bin/env python
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08002# -*- 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 Huanga0e55b82016-02-10 14:32:07 +080010import ctypes
11import ctypes.util
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
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +080017import platform
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080018import Queue
Wei-Ning Huang829e0c82015-05-26 14:37:23 +080019import re
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080020import select
Wei-Ning Huanga301f572015-06-03 17:34:21 +080021import signal
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080022import socket
Wei-Ning Huangf5311a02016-02-04 15:23:46 +080023import ssl
Wei-Ning Huangb05cde32015-08-01 09:48:41 +080024import struct
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080025import subprocess
26import sys
Moja Hsuc9ecc8b2015-07-13 11:39:17 +080027import termios
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080028import threading
29import time
Joel Kitching22b89042015-08-06 18:23:29 +080030import traceback
Wei-Ning Huang39169902015-09-19 06:00:23 +080031import tty
Wei-Ning Huange0def6a2015-11-05 15:41:24 +080032import urllib2
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080033import uuid
34
Wei-Ning Huang2132de32015-04-13 17:24:38 +080035import jsonrpclib
36from jsonrpclib.SimpleJSONRPCServer import SimpleJSONRPCServer
37
38
Peter Shihc56c5b62016-12-22 12:30:57 +080039_GHOST_RPC_PORT = int(os.getenv('GHOST_RPC_PORT', 4499))
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080040
Peter Shihc56c5b62016-12-22 12:30:57 +080041_OVERLORD_PORT = int(os.getenv('OVERLORD_PORT', 4455))
42_OVERLORD_LAN_DISCOVERY_PORT = int(os.getenv('OVERLORD_LD_PORT', 4456))
43_OVERLORD_HTTP_PORT = int(os.getenv('OVERLORD_HTTP_PORT', 9000))
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080044
45_BUFSIZE = 8192
46_RETRY_INTERVAL = 2
47_SEPARATOR = '\r\n'
48_PING_TIMEOUT = 3
49_PING_INTERVAL = 5
50_REQUEST_TIMEOUT_SECS = 60
51_SHELL = os.getenv('SHELL', '/bin/bash')
Wei-Ning Huang7dbf4a72016-03-02 20:16:20 +080052_DEFAULT_BIND_ADDRESS = 'localhost'
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080053
Moja Hsuc9ecc8b2015-07-13 11:39:17 +080054_CONTROL_START = 128
55_CONTROL_END = 129
56
Wei-Ning Huanga301f572015-06-03 17:34:21 +080057_BLOCK_SIZE = 4096
Wei-Ning Huange0def6a2015-11-05 15:41:24 +080058_CONNECT_TIMEOUT = 3
Wei-Ning Huanga301f572015-06-03 17:34:21 +080059
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +080060# Stream control
61_STDIN_CLOSED = '##STDIN_CLOSED##'
62
Wei-Ning Huang7ec55342015-09-17 08:46:06 +080063SUCCESS = 'success'
64FAILED = 'failed'
65DISCONNECTED = 'disconnected'
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080066
Joel Kitching22b89042015-08-06 18:23:29 +080067
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080068class PingTimeoutError(Exception):
69 pass
70
71
72class RequestError(Exception):
73 pass
74
75
Wei-Ning Huangf5311a02016-02-04 15:23:46 +080076class BufferedSocket(object):
Wei-Ning Huanga28cd232016-01-27 15:04:41 +080077 """A buffered socket that supports unrecv.
78
79 Allow putting back data back to the socket for the next recv() call.
80 """
Wei-Ning Huangf5311a02016-02-04 15:23:46 +080081 def __init__(self, sock):
82 self.sock = sock
Wei-Ning Huanga28cd232016-01-27 15:04:41 +080083 self._buf = ''
84
Wei-Ning Huangf5311a02016-02-04 15:23:46 +080085 def fileno(self):
86 return self.sock.fileno()
87
Wei-Ning Huanga28cd232016-01-27 15:04:41 +080088 def Recv(self, bufsize, flags=0):
89 if self._buf:
90 if len(self._buf) >= bufsize:
91 ret = self._buf[:bufsize]
92 self._buf = self._buf[bufsize:]
93 return ret
94 else:
95 ret = self._buf
96 self._buf = ''
Wei-Ning Huangf5311a02016-02-04 15:23:46 +080097 return ret + self.sock.recv(bufsize - len(ret), flags)
Wei-Ning Huanga28cd232016-01-27 15:04:41 +080098 else:
Wei-Ning Huangf5311a02016-02-04 15:23:46 +080099 return self.sock.recv(bufsize, flags)
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800100
101 def UnRecv(self, buf):
102 self._buf = buf + self._buf
103
104 def Send(self, *args, **kwargs):
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800105 return self.sock.send(*args, **kwargs)
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800106
107 def RecvBuf(self):
108 """Only recive from buffer."""
109 ret = self._buf
110 self._buf = ''
111 return ret
112
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800113 def Close(self):
114 self.sock.close()
115
116
117class TLSSettings(object):
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800118 def __init__(self, tls_cert_file, verify):
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800119 """Constructor.
120
121 Args:
122 tls_cert_file: TLS certificate in PEM format.
123 enable_tls_without_verify: enable TLS but don't verify certificate.
124 """
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800125 self._enabled = False
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800126 self._tls_cert_file = tls_cert_file
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800127 self._verify = verify
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800128 self._tls_context = None
129
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800130 def _UpdateContext(self):
131 if not self._enabled:
132 self._tls_context = None
133 return
134
135 self._tls_context = ssl.SSLContext(ssl.PROTOCOL_TLSv1_2)
136 self._tls_context.verify_mode = ssl.CERT_REQUIRED
137
138 if self._verify:
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800139 if self._tls_cert_file:
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800140 self._tls_context.check_hostname = True
141 try:
142 self._tls_context.load_verify_locations(self._tls_cert_file)
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800143 logging.info('TLSSettings: using user-supplied ca-certificate')
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800144 except IOError as e:
145 logging.error('TLSSettings: %s: %s', self._tls_cert_file, e)
146 sys.exit(1)
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800147 else:
148 self._tls_context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
149 logging.info('TLSSettings: using built-in ca-certificates')
150 else:
151 self._tls_context.verify_mode = ssl.CERT_NONE
152 logging.info('TLSSettings: skipping TLS verification!!!')
153
154 def SetEnabled(self, enabled):
155 logging.info('TLSSettings: enabled: %s', enabled)
156
157 if self._enabled != enabled:
158 self._enabled = enabled
159 self._UpdateContext()
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800160
161 def Enabled(self):
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800162 return self._enabled
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800163
164 def Context(self):
165 return self._tls_context
166
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800167
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800168class Ghost(object):
169 """Ghost implements the client protocol of Overlord.
170
171 Ghost provide terminal/shell/logcat functionality and manages the client
172 side connectivity.
173 """
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800174 NONE, AGENT, TERMINAL, SHELL, LOGCAT, FILE, FORWARD = range(7)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800175
176 MODE_NAME = {
177 NONE: 'NONE',
178 AGENT: 'Agent',
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800179 TERMINAL: 'Terminal',
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800180 SHELL: 'Shell',
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800181 LOGCAT: 'Logcat',
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800182 FILE: 'File',
183 FORWARD: 'Forward'
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800184 }
185
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +0800186 RANDOM_MID = '##random_mid##'
187
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800188 def __init__(self, overlord_addrs, tls_settings=None, mode=AGENT, mid=None,
189 sid=None, prop_file=None, terminal_sid=None, tty_device=None,
Peter Shih220a96d2016-12-22 17:02:16 +0800190 command=None, file_op=None, port=None, tls_mode=None):
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800191 """Constructor.
192
193 Args:
194 overlord_addrs: a list of possible address of overlord.
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800195 tls_settings: a TLSSetting object.
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800196 mode: client mode, either AGENT, SHELL or LOGCAT
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +0800197 mid: a str to set for machine ID. If mid equals Ghost.RANDOM_MID, machine
198 id is randomly generated.
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800199 sid: session ID. If the connection is requested by overlord, sid should
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800200 be set to the corresponding session id assigned by overlord.
Wei-Ning Huang23ed0162015-09-18 14:42:03 +0800201 prop_file: properties file filename.
Wei-Ning Huangd521f282015-08-07 05:28:04 +0800202 terminal_sid: the terminal session ID associate with this client. This is
203 use for file download.
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800204 tty_device: the terminal device to open, if tty_device is None, as pseudo
205 terminal will be opened instead.
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800206 command: the command to execute when we are in SHELL mode.
Wei-Ning Huang8ee3bcd2015-10-01 17:10:01 +0800207 file_op: a tuple (action, filepath, perm). action is either 'download' or
208 'upload'. perm is the permission to set for the file.
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800209 port: port number to forward.
Peter Shih220a96d2016-12-22 17:02:16 +0800210 tls_mode: can be [True, False, None]. if not None, skip detection of
211 TLS and assume whether server use TLS or not.
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800212 """
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800213 assert mode in [Ghost.AGENT, Ghost.TERMINAL, Ghost.SHELL, Ghost.FILE,
214 Ghost.FORWARD]
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800215 if mode == Ghost.SHELL:
216 assert command is not None
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800217 if mode == Ghost.FILE:
218 assert file_op is not None
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800219
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800220 self._platform = platform.system()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800221 self._overlord_addrs = overlord_addrs
Wei-Ning Huangad330c52015-03-12 20:34:18 +0800222 self._connected_addr = None
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800223 self._tls_settings = tls_settings
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +0800224 self._mid = mid
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800225 self._sock = None
Wei-Ning Huang7ec55342015-09-17 08:46:06 +0800226 self._mode = mode
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800227 self._machine_id = self.GetMachineID()
Wei-Ning Huangfed95862015-08-07 03:17:11 +0800228 self._session_id = sid if sid is not None else str(uuid.uuid4())
Wei-Ning Huangd521f282015-08-07 05:28:04 +0800229 self._terminal_session_id = terminal_sid
Wei-Ning Huang23ed0162015-09-18 14:42:03 +0800230 self._ttyname_to_sid = {}
231 self._terminal_sid_to_pid = {}
232 self._prop_file = prop_file
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800233 self._properties = {}
Wei-Ning Huang7ec55342015-09-17 08:46:06 +0800234 self._register_status = DISCONNECTED
235 self._reset = threading.Event()
Peter Shih220a96d2016-12-22 17:02:16 +0800236 self._tls_mode = tls_mode
Wei-Ning Huang7ec55342015-09-17 08:46:06 +0800237
238 # RPC
Wei-Ning Huang7ec55342015-09-17 08:46:06 +0800239 self._requests = {}
240 self._queue = Queue.Queue()
241
242 # Protocol specific
243 self._last_ping = 0
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800244 self._tty_device = tty_device
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800245 self._shell_command = command
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800246 self._file_op = file_op
Wei-Ning Huang7ec55342015-09-17 08:46:06 +0800247 self._download_queue = Queue.Queue()
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800248 self._port = port
Wei-Ning Huang7ec55342015-09-17 08:46:06 +0800249
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800250 def SetIgnoreChild(self, status):
251 # Only ignore child for Agent since only it could spawn child Ghost.
252 if self._mode == Ghost.AGENT:
253 signal.signal(signal.SIGCHLD,
254 signal.SIG_IGN if status else signal.SIG_DFL)
255
256 def GetFileSha1(self, filename):
257 with open(filename, 'r') as f:
258 return hashlib.sha1(f.read()).hexdigest()
259
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800260 def TLSEnabled(self, host, port):
261 """Determine if TLS is enabled on given server address."""
Wei-Ning Huang58833882015-09-16 16:52:37 +0800262 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
263 try:
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800264 # Allow any certificate since we only want to check if server talks TLS.
265 context = ssl.SSLContext(ssl.PROTOCOL_TLSv1_2)
266 context.verify_mode = ssl.CERT_NONE
Wei-Ning Huang58833882015-09-16 16:52:37 +0800267
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800268 sock = context.wrap_socket(sock, server_hostname=host)
269 sock.settimeout(_CONNECT_TIMEOUT)
270 sock.connect((host, port))
271 return True
Wei-Ning Huangb6605d22016-06-22 17:33:37 +0800272 except ssl.SSLError:
273 return False
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800274 except socket.error: # Connect refused or timeout
275 raise
Wei-Ning Huang58833882015-09-16 16:52:37 +0800276 except Exception:
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800277 return False # For whatever reason above failed, assume False
Wei-Ning Huang58833882015-09-16 16:52:37 +0800278
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800279 def Upgrade(self):
280 logging.info('Upgrade: initiating upgrade sequence...')
281
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800282 try:
Wei-Ning Huangb6605d22016-06-22 17:33:37 +0800283 https_enabled = self.TLSEnabled(self._connected_addr[0],
284 _OVERLORD_HTTP_PORT)
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800285 except socket.error:
Wei-Ning Huangb6605d22016-06-22 17:33:37 +0800286 logging.error('Upgrade: failed to connect to Overlord HTTP server, '
287 'abort')
288 return
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800289
290 if self._tls_settings.Enabled() and not https_enabled:
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800291 logging.error('Upgrade: TLS enforced but found Overlord HTTP server '
292 'without TLS enabled! Possible mis-configuration or '
293 'DNS/IP spoofing detected, abort')
294 return
295
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800296 scriptpath = os.path.abspath(sys.argv[0])
Wei-Ning Huang03f9f762015-09-16 21:51:35 +0800297 url = 'http%s://%s:%d/upgrade/ghost.py' % (
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800298 's' if https_enabled else '', self._connected_addr[0],
Wei-Ning Huang03f9f762015-09-16 21:51:35 +0800299 _OVERLORD_HTTP_PORT)
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800300
301 # Download sha1sum for ghost.py for verification
302 try:
Wei-Ning Huange0def6a2015-11-05 15:41:24 +0800303 with contextlib.closing(
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800304 urllib2.urlopen(url + '.sha1', timeout=_CONNECT_TIMEOUT,
305 context=self._tls_settings.Context())) as f:
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800306 if f.getcode() != 200:
307 raise RuntimeError('HTTP status %d' % f.getcode())
308 sha1sum = f.read().strip()
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800309 except (ssl.SSLError, ssl.CertificateError) as e:
310 logging.error('Upgrade: %s: %s', e.__class__.__name__, e)
311 return
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800312 except Exception:
313 logging.error('Upgrade: failed to download sha1sum file, abort')
314 return
315
316 if self.GetFileSha1(scriptpath) == sha1sum:
317 logging.info('Upgrade: ghost is already up-to-date, skipping upgrade')
318 return
319
320 # Download upgrade version of ghost.py
321 try:
Wei-Ning Huange0def6a2015-11-05 15:41:24 +0800322 with contextlib.closing(
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800323 urllib2.urlopen(url, timeout=_CONNECT_TIMEOUT,
324 context=self._tls_settings.Context())) as f:
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800325 if f.getcode() != 200:
326 raise RuntimeError('HTTP status %d' % f.getcode())
327 data = f.read()
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800328 except (ssl.SSLError, ssl.CertificateError) as e:
329 logging.error('Upgrade: %s: %s', e.__class__.__name__, e)
330 return
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800331 except Exception:
332 logging.error('Upgrade: failed to download upgrade, abort')
333 return
334
335 # Compare SHA1 sum
336 if hashlib.sha1(data).hexdigest() != sha1sum:
337 logging.error('Upgrade: sha1sum mismatch, abort')
338 return
339
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800340 try:
341 with open(scriptpath, 'w') as f:
342 f.write(data)
343 except Exception:
344 logging.error('Upgrade: failed to write upgrade onto disk, abort')
345 return
346
347 logging.info('Upgrade: restarting ghost...')
348 self.CloseSockets()
349 self.SetIgnoreChild(False)
Wei-Ning Huang5f3fa8f2015-10-24 15:08:48 +0800350 os.execve(scriptpath, [scriptpath] + sys.argv[1:], os.environ)
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800351
Wei-Ning Huang23ed0162015-09-18 14:42:03 +0800352 def LoadProperties(self):
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800353 try:
Wei-Ning Huang23ed0162015-09-18 14:42:03 +0800354 if self._prop_file:
355 with open(self._prop_file, 'r') as f:
356 self._properties = json.loads(f.read())
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800357 except Exception as e:
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800358 logging.error('LoadProperties: ' + str(e))
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800359
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800360 def CloseSockets(self):
361 # Close sockets opened by parent process, since we don't use it anymore.
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800362 if self._platform == 'Linux':
363 for fd in os.listdir('/proc/self/fd/'):
364 try:
365 real_fd = os.readlink('/proc/self/fd/%s' % fd)
366 if real_fd.startswith('socket'):
367 os.close(int(fd))
368 except Exception:
369 pass
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800370
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800371 def SpawnGhost(self, mode, sid=None, terminal_sid=None, tty_device=None,
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800372 command=None, file_op=None, port=None):
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800373 """Spawn a child ghost with specific mode.
374
375 Returns:
376 The spawned child process pid.
377 """
Joel Kitching22b89042015-08-06 18:23:29 +0800378 # Restore the default signal handler, so our child won't have problems.
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800379 self.SetIgnoreChild(False)
380
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800381 pid = os.fork()
382 if pid == 0:
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800383 self.CloseSockets()
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800384 g = Ghost([self._connected_addr], tls_settings=self._tls_settings,
385 mode=mode, mid=Ghost.RANDOM_MID, sid=sid,
Wei-Ning Huang23ed0162015-09-18 14:42:03 +0800386 terminal_sid=terminal_sid, tty_device=tty_device,
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800387 command=command, file_op=file_op, port=port)
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800388 g.Start()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800389 sys.exit(0)
390 else:
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800391 self.SetIgnoreChild(True)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800392 return pid
393
394 def Timestamp(self):
395 return int(time.time())
396
397 def GetGateWayIP(self):
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800398 if self._platform == 'Darwin':
399 output = subprocess.check_output(['route', '-n', 'get', 'default'])
400 ret = re.search('gateway: (.*)', output)
401 if ret:
402 return [ret.group(1)]
403 elif self._platform == 'Linux':
404 with open('/proc/net/route', 'r') as f:
405 lines = f.readlines()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800406
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800407 ips = []
408 for line in lines:
409 parts = line.split('\t')
410 if parts[2] == '00000000':
411 continue
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800412
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800413 try:
414 h = parts[2].decode('hex')
415 ips.append('%d.%d.%d.%d' % tuple(ord(x) for x in reversed(h)))
416 except TypeError:
417 pass
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800418
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800419 return ips
420 else:
421 logging.warning('GetGateWayIP: unsupported platform')
422 return []
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800423
Wei-Ning Huang829e0c82015-05-26 14:37:23 +0800424 def GetShopfloorIP(self):
425 try:
426 import factory_common # pylint: disable=W0612
427 from cros.factory.test import shopfloor
428
429 url = shopfloor.get_server_url()
430 match = re.match(r'^https?://(.*):.*$', url)
431 if match:
432 return [match.group(1)]
433 except Exception:
434 pass
435 return []
436
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800437 def GetMachineID(self):
438 """Generates machine-dependent ID string for a machine.
439 There are many ways to generate a machine ID:
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800440 Linux:
441 1. factory device_id
442 2. factory device-data
443 3. /sys/class/dmi/id/product_uuid (only available on intel machines)
444 4. MAC address
445 We follow the listed order to generate machine ID, and fallback to the
446 next alternative if the previous doesn't work.
447
448 Darwin:
449 All Darwin system should have the IOPlatformSerialNumber attribute.
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800450 """
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +0800451 if self._mid == Ghost.RANDOM_MID:
Wei-Ning Huangaed90452015-03-23 17:50:21 +0800452 return str(uuid.uuid4())
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +0800453 elif self._mid:
454 return self._mid
Wei-Ning Huangaed90452015-03-23 17:50:21 +0800455
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800456 # Darwin
457 if self._platform == 'Darwin':
458 output = subprocess.check_output(['ioreg', '-rd1', '-c',
459 'IOPlatformExpertDevice'])
460 ret = re.search('"IOPlatformSerialNumber" = "(.*)"', output)
461 if ret:
462 return ret.group(1)
463
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800464 # Try factory device id
465 try:
466 import factory_common # pylint: disable=W0612
467 from cros.factory.test import event_log
468 with open(event_log.DEVICE_ID_PATH) as f:
469 return f.read().strip()
470 except Exception:
471 pass
472
Wei-Ning Huang1d7603b2015-07-03 17:38:56 +0800473 # Try factory device data
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800474 try:
475 p = subprocess.Popen('factory device-data | grep mlb_serial_number | '
476 'cut -d " " -f 2', stdout=subprocess.PIPE,
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800477 stderr=subprocess.PIPE, shell=True)
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800478 stdout, unused_stderr = p.communicate()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800479 if stdout == '':
Wei-Ning Huang1d7603b2015-07-03 17:38:56 +0800480 raise RuntimeError('empty mlb number')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800481 return stdout.strip()
482 except Exception:
483 pass
484
Wei-Ning Huang1d7603b2015-07-03 17:38:56 +0800485 # Try DMI product UUID
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800486 try:
487 with open('/sys/class/dmi/id/product_uuid', 'r') as f:
488 return f.read().strip()
489 except Exception:
490 pass
491
Wei-Ning Huang1d7603b2015-07-03 17:38:56 +0800492 # Use MAC address if non is available
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800493 try:
494 macs = []
495 ifaces = sorted(os.listdir('/sys/class/net'))
496 for iface in ifaces:
497 if iface == 'lo':
498 continue
499
500 with open('/sys/class/net/%s/address' % iface, 'r') as f:
501 macs.append(f.read().strip())
502
503 return ';'.join(macs)
504 except Exception:
505 pass
506
Moja Hsuc9ecc8b2015-07-13 11:39:17 +0800507 raise RuntimeError('can\'t generate machine ID')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800508
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800509 def GetProcessWorkingDirectory(self, pid):
510 if self._platform == 'Linux':
511 return os.readlink('/proc/%d/cwd' % pid)
512 elif self._platform == 'Darwin':
513 PROC_PIDVNODEPATHINFO = 9
514 proc_vnodepathinfo_size = 2352
515 vid_path_offset = 152
516
517 proc = ctypes.cdll.LoadLibrary(ctypes.util.find_library('libproc'))
518 buf = ctypes.create_string_buffer('\0' * proc_vnodepathinfo_size)
519 proc.proc_pidinfo(pid, PROC_PIDVNODEPATHINFO, 0,
520 ctypes.byref(buf), proc_vnodepathinfo_size)
521 buf = buf.raw[vid_path_offset:]
522 n = buf.index('\0')
523 return buf[:n]
524 else:
525 raise RuntimeError('GetProcessWorkingDirectory: unsupported platform')
526
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800527 def Reset(self):
528 """Reset state and clear request handlers."""
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800529 if self._sock is not None:
530 self._sock.Close()
531 self._sock = None
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800532 self._reset.clear()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800533 self._last_ping = 0
534 self._requests = {}
Wei-Ning Huang23ed0162015-09-18 14:42:03 +0800535 self.LoadProperties()
Wei-Ning Huang7ec55342015-09-17 08:46:06 +0800536 self._register_status = DISCONNECTED
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800537
538 def SendMessage(self, msg):
539 """Serialize the message and send it through the socket."""
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800540 self._sock.Send(json.dumps(msg) + _SEPARATOR)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800541
542 def SendRequest(self, name, args, handler=None,
543 timeout=_REQUEST_TIMEOUT_SECS):
544 if handler and not callable(handler):
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800545 raise RequestError('Invalid request handler for msg "%s"' % name)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800546
547 rid = str(uuid.uuid4())
548 msg = {'rid': rid, 'timeout': timeout, 'name': name, 'params': args}
Wei-Ning Huange2981862015-08-03 15:03:08 +0800549 if timeout >= 0:
550 self._requests[rid] = [self.Timestamp(), timeout, handler]
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800551 self.SendMessage(msg)
552
553 def SendResponse(self, omsg, status, params=None):
554 msg = {'rid': omsg['rid'], 'response': status, 'params': params}
555 self.SendMessage(msg)
556
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800557 def HandleTTYControl(self, fd, control_str):
558 msg = json.loads(control_str)
Moja Hsuc9ecc8b2015-07-13 11:39:17 +0800559 command = msg['command']
560 params = msg['params']
561 if command == 'resize':
562 # some error happened on websocket
563 if len(params) != 2:
564 return
565 winsize = struct.pack('HHHH', params[0], params[1], 0, 0)
566 fcntl.ioctl(fd, termios.TIOCSWINSZ, winsize)
567 else:
568 logging.warn('Invalid request command "%s"', command)
569
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800570 def SpawnTTYServer(self, unused_var):
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800571 """Spawn a TTY server and forward I/O to the TCP socket."""
572 logging.info('SpawnTTYServer: started')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800573
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800574 try:
575 if self._tty_device is None:
576 pid, fd = os.forkpty()
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800577
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800578 if pid == 0:
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800579 ttyname = os.ttyname(sys.stdout.fileno())
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800580 try:
581 server = GhostRPCServer()
582 server.RegisterTTY(self._session_id, ttyname)
583 server.RegisterSession(self._session_id, os.getpid())
584 except Exception:
585 # If ghost is launched without RPC server, the call will fail but we
586 # can ignore it.
587 pass
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800588
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800589 # The directory that contains the current running ghost script
590 script_dir = os.path.dirname(os.path.abspath(sys.argv[0]))
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800591
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800592 env = os.environ.copy()
593 env['USER'] = os.getenv('USER', 'root')
594 env['HOME'] = os.getenv('HOME', '/root')
595 env['PATH'] = os.getenv('PATH') + ':%s' % script_dir
596 os.chdir(env['HOME'])
597 os.execve(_SHELL, [_SHELL], env)
598 else:
599 fd = os.open(self._tty_device, os.O_RDWR)
Wei-Ning Huang39169902015-09-19 06:00:23 +0800600 tty.setraw(fd)
601 attr = termios.tcgetattr(fd)
602 attr[0] &= ~(termios.IXON | termios.IXOFF)
603 attr[2] |= termios.CLOCAL
604 attr[2] &= ~termios.CRTSCTS
605 attr[4] = termios.B115200
606 attr[5] = termios.B115200
607 termios.tcsetattr(fd, termios.TCSANOW, attr)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800608
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800609 nonlocals = {'control_state': None, 'control_str': ''}
610
611 def _ProcessBuffer(buf):
612 write_buffer = ''
613 while buf:
614 if nonlocals['control_state']:
615 if chr(_CONTROL_END) in buf:
616 index = buf.index(chr(_CONTROL_END))
617 nonlocals['control_str'] += buf[:index]
618 self.HandleTTYControl(fd, nonlocals['control_str'])
619 nonlocals['control_state'] = None
620 nonlocals['control_str'] = ''
621 buf = buf[index+1:]
622 else:
623 nonlocals['control_str'] += buf
624 buf = ''
625 else:
626 if chr(_CONTROL_START) in buf:
627 nonlocals['control_state'] = _CONTROL_START
628 index = buf.index(chr(_CONTROL_START))
629 write_buffer += buf[:index]
630 buf = buf[index+1:]
631 else:
632 write_buffer += buf
633 buf = ''
634
635 if write_buffer:
636 os.write(fd, write_buffer)
637
638 _ProcessBuffer(self._sock.RecvBuf())
639
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800640 while True:
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800641 rd, unused_wd, unused_xd = select.select([self._sock, fd], [], [])
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800642
643 if fd in rd:
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800644 self._sock.Send(os.read(fd, _BUFSIZE))
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800645
646 if self._sock in rd:
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800647 buf = self._sock.Recv(_BUFSIZE)
648 if len(buf) == 0:
Wei-Ning Huang5f3fa8f2015-10-24 15:08:48 +0800649 raise RuntimeError('connection terminated')
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800650 _ProcessBuffer(buf)
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800651 except Exception as e:
652 logging.error('SpawnTTYServer: %s', e)
653 finally:
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800654 self._sock.Close()
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800655
656 logging.info('SpawnTTYServer: terminated')
657 sys.exit(0)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800658
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800659 def SpawnShellServer(self, unused_var):
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800660 """Spawn a shell server and forward input/output from/to the TCP socket."""
661 logging.info('SpawnShellServer: started')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800662
Wei-Ning Huang5f3fa8f2015-10-24 15:08:48 +0800663 # Add ghost executable to PATH
664 script_dir = os.path.dirname(os.path.abspath(sys.argv[0]))
665 env = os.environ.copy()
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800666 env['PATH'] = '%s:%s' % (script_dir, os.getenv('PATH'))
Wei-Ning Huang5f3fa8f2015-10-24 15:08:48 +0800667
668 # Execute shell command from HOME directory
669 os.chdir(os.getenv('HOME', '/tmp'))
670
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800671 p = subprocess.Popen(self._shell_command, stdin=subprocess.PIPE,
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800672 stdout=subprocess.PIPE, stderr=subprocess.PIPE,
Wei-Ning Huang5f3fa8f2015-10-24 15:08:48 +0800673 shell=True, env=env)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800674
675 def make_non_block(fd):
676 fl = fcntl.fcntl(fd, fcntl.F_GETFL)
677 fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
678
679 make_non_block(p.stdout)
680 make_non_block(p.stderr)
681
682 try:
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800683 p.stdin.write(self._sock.RecvBuf())
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800684
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800685 while True:
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800686 rd, unused_wd, unused_xd = select.select(
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800687 [p.stdout, p.stderr, self._sock], [], [])
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800688 if p.stdout in rd:
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800689 self._sock.Send(p.stdout.read(_BUFSIZE))
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800690
691 if p.stderr in rd:
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800692 self._sock.Send(p.stderr.read(_BUFSIZE))
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800693
694 if self._sock in rd:
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800695 ret = self._sock.Recv(_BUFSIZE)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800696 if len(ret) == 0:
Wei-Ning Huang5f3fa8f2015-10-24 15:08:48 +0800697 raise RuntimeError('connection terminated')
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800698
699 try:
700 idx = ret.index(_STDIN_CLOSED * 2)
701 p.stdin.write(ret[:idx])
702 p.stdin.close()
703 except ValueError:
704 p.stdin.write(ret)
Wei-Ning Huangf14c84e2015-08-03 15:03:08 +0800705 p.poll()
706 if p.returncode != None:
707 break
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800708 except Exception as e:
709 logging.error('SpawnShellServer: %s', e)
Wei-Ning Huangf14c84e2015-08-03 15:03:08 +0800710 finally:
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800711 # Check if the process is terminated. If not, Send SIGTERM to process,
712 # then wait for 1 second. Send another SIGKILL to make sure the process is
713 # terminated.
714 p.poll()
715 if p.returncode is None:
716 try:
717 p.terminate()
718 time.sleep(1)
719 p.kill()
720 except Exception:
721 pass
Wei-Ning Huang5f3fa8f2015-10-24 15:08:48 +0800722
723 p.wait()
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800724 self._sock.Close()
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800725
726 logging.info('SpawnShellServer: terminated')
727 sys.exit(0)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800728
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800729 def InitiateFileOperation(self, unused_var):
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800730 if self._file_op[0] == 'download':
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800731 try:
732 size = os.stat(self._file_op[1]).st_size
733 except OSError as e:
734 logging.error('InitiateFileOperation: download: %s', e)
735 sys.exit(1)
736
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800737 self.SendRequest('request_to_download',
Wei-Ning Huangd521f282015-08-07 05:28:04 +0800738 {'terminal_sid': self._terminal_session_id,
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800739 'filename': os.path.basename(self._file_op[1]),
740 'size': size})
Wei-Ning Huange2981862015-08-03 15:03:08 +0800741 elif self._file_op[0] == 'upload':
742 self.SendRequest('clear_to_upload', {}, timeout=-1)
743 self.StartUploadServer()
744 else:
745 logging.error('InitiateFileOperation: unknown file operation, ignored')
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800746
747 def StartDownloadServer(self):
748 logging.info('StartDownloadServer: started')
749
750 try:
751 with open(self._file_op[1], 'rb') as f:
752 while True:
753 data = f.read(_BLOCK_SIZE)
754 if len(data) == 0:
755 break
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800756 self._sock.Send(data)
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800757 except Exception as e:
758 logging.error('StartDownloadServer: %s', e)
759 finally:
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800760 self._sock.Close()
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800761
762 logging.info('StartDownloadServer: terminated')
763 sys.exit(0)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800764
Wei-Ning Huange2981862015-08-03 15:03:08 +0800765 def StartUploadServer(self):
766 logging.info('StartUploadServer: started')
Wei-Ning Huange2981862015-08-03 15:03:08 +0800767 try:
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800768 filepath = self._file_op[1]
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800769 dirname = os.path.dirname(filepath)
770 if not os.path.exists(dirname):
771 try:
772 os.makedirs(dirname)
773 except Exception:
774 pass
Wei-Ning Huange2981862015-08-03 15:03:08 +0800775
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800776 with open(filepath, 'wb') as f:
Wei-Ning Huang8ee3bcd2015-10-01 17:10:01 +0800777 if self._file_op[2]:
778 os.fchmod(f.fileno(), self._file_op[2])
779
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800780 f.write(self._sock.RecvBuf())
781
Wei-Ning Huange2981862015-08-03 15:03:08 +0800782 while True:
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800783 rd, unused_wd, unused_xd = select.select([self._sock], [], [])
Wei-Ning Huange2981862015-08-03 15:03:08 +0800784 if self._sock in rd:
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800785 buf = self._sock.Recv(_BLOCK_SIZE)
Wei-Ning Huange2981862015-08-03 15:03:08 +0800786 if len(buf) == 0:
787 break
788 f.write(buf)
789 except socket.error as e:
790 logging.error('StartUploadServer: socket error: %s', e)
791 except Exception as e:
792 logging.error('StartUploadServer: %s', e)
793 finally:
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800794 self._sock.Close()
Wei-Ning Huange2981862015-08-03 15:03:08 +0800795
796 logging.info('StartUploadServer: terminated')
797 sys.exit(0)
798
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800799 def SpawnPortForwardServer(self, unused_var):
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800800 """Spawn a port forwarding server and forward I/O to the TCP socket."""
801 logging.info('SpawnPortForwardServer: started')
802
803 src_sock = None
804 try:
805 src_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
Wei-Ning Huange0def6a2015-11-05 15:41:24 +0800806 src_sock.settimeout(_CONNECT_TIMEOUT)
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800807 src_sock.connect(('localhost', self._port))
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800808
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800809 src_sock.send(self._sock.RecvBuf())
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800810
811 while True:
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800812 rd, unused_wd, unused_xd = select.select([self._sock, src_sock], [], [])
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800813
814 if self._sock in rd:
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800815 data = self._sock.Recv(_BUFSIZE)
Wei-Ning Huang5f3fa8f2015-10-24 15:08:48 +0800816 if len(data) == 0:
817 raise RuntimeError('connection terminated')
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800818 src_sock.send(data)
819
820 if src_sock in rd:
821 data = src_sock.recv(_BUFSIZE)
Wei-Ning Huang5f3fa8f2015-10-24 15:08:48 +0800822 if len(data) == 0:
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800823 break
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800824 self._sock.Send(data)
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800825 except Exception as e:
826 logging.error('SpawnPortForwardServer: %s', e)
827 finally:
828 if src_sock:
829 src_sock.close()
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800830 self._sock.Close()
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800831
832 logging.info('SpawnPortForwardServer: terminated')
833 sys.exit(0)
834
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800835 def Ping(self):
836 def timeout_handler(x):
837 if x is None:
838 raise PingTimeoutError
839
840 self._last_ping = self.Timestamp()
841 self.SendRequest('ping', {}, timeout_handler, 5)
842
Wei-Ning Huangae923642015-09-24 14:08:09 +0800843 def HandleFileDownloadRequest(self, msg):
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800844 params = msg['params']
Wei-Ning Huangae923642015-09-24 14:08:09 +0800845 filepath = params['filename']
846 if not os.path.isabs(filepath):
847 filepath = os.path.join(os.getenv('HOME', '/tmp'), filepath)
848
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800849 try:
Wei-Ning Huang11c35022015-10-21 16:52:32 +0800850 with open(filepath, 'r') as _:
Wei-Ning Huang46a3fc92015-10-06 02:35:27 +0800851 pass
852 except Exception as e:
Wei-Ning Huangae923642015-09-24 14:08:09 +0800853 return self.SendResponse(msg, str(e))
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800854
855 self.SpawnGhost(self.FILE, params['sid'],
Wei-Ning Huangae923642015-09-24 14:08:09 +0800856 file_op=('download', filepath))
857 self.SendResponse(msg, SUCCESS)
858
859 def HandleFileUploadRequest(self, msg):
860 params = msg['params']
861
862 # Resolve upload filepath
863 filename = params['filename']
864 dest_path = filename
865
866 # If dest is specified, use it first
867 dest_path = params.get('dest', '')
868 if dest_path:
869 if not os.path.isabs(dest_path):
870 dest_path = os.path.join(os.getenv('HOME', '/tmp'), dest_path)
871
872 if os.path.isdir(dest_path):
873 dest_path = os.path.join(dest_path, filename)
874 else:
875 target_dir = os.getenv('HOME', '/tmp')
876
877 # Terminal session ID found, upload to it's current working directory
878 if params.has_key('terminal_sid'):
879 pid = self._terminal_sid_to_pid.get(params['terminal_sid'], None)
880 if pid:
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800881 try:
882 target_dir = self.GetProcessWorkingDirectory(pid)
883 except Exception as e:
884 logging.error(e)
Wei-Ning Huangae923642015-09-24 14:08:09 +0800885
886 dest_path = os.path.join(target_dir, filename)
887
888 try:
889 os.makedirs(os.path.dirname(dest_path))
890 except Exception:
891 pass
892
893 try:
894 with open(dest_path, 'w') as _:
895 pass
896 except Exception as e:
897 return self.SendResponse(msg, str(e))
898
Wei-Ning Huangd6f69762015-10-01 21:02:07 +0800899 # If not check_only, spawn FILE mode ghost agent to handle upload
900 if not params.get('check_only', False):
901 self.SpawnGhost(self.FILE, params['sid'],
902 file_op=('upload', dest_path, params.get('perm', None)))
Wei-Ning Huang7ec55342015-09-17 08:46:06 +0800903 self.SendResponse(msg, SUCCESS)
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800904
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800905 def HandleRequest(self, msg):
Wei-Ning Huange2981862015-08-03 15:03:08 +0800906 command = msg['name']
907 params = msg['params']
908
909 if command == 'upgrade':
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800910 self.Upgrade()
Wei-Ning Huange2981862015-08-03 15:03:08 +0800911 elif command == 'terminal':
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800912 self.SpawnGhost(self.TERMINAL, params['sid'],
913 tty_device=params['tty_device'])
Wei-Ning Huang7ec55342015-09-17 08:46:06 +0800914 self.SendResponse(msg, SUCCESS)
Wei-Ning Huange2981862015-08-03 15:03:08 +0800915 elif command == 'shell':
916 self.SpawnGhost(self.SHELL, params['sid'], command=params['command'])
Wei-Ning Huang7ec55342015-09-17 08:46:06 +0800917 self.SendResponse(msg, SUCCESS)
Wei-Ning Huange2981862015-08-03 15:03:08 +0800918 elif command == 'file_download':
Wei-Ning Huangae923642015-09-24 14:08:09 +0800919 self.HandleFileDownloadRequest(msg)
Wei-Ning Huange2981862015-08-03 15:03:08 +0800920 elif command == 'clear_to_download':
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800921 self.StartDownloadServer()
Wei-Ning Huange2981862015-08-03 15:03:08 +0800922 elif command == 'file_upload':
Wei-Ning Huangae923642015-09-24 14:08:09 +0800923 self.HandleFileUploadRequest(msg)
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800924 elif command == 'forward':
925 self.SpawnGhost(self.FORWARD, params['sid'], port=params['port'])
926 self.SendResponse(msg, SUCCESS)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800927
928 def HandleResponse(self, response):
929 rid = str(response['rid'])
930 if rid in self._requests:
931 handler = self._requests[rid][2]
932 del self._requests[rid]
933 if callable(handler):
934 handler(response)
935 else:
Joel Kitching22b89042015-08-06 18:23:29 +0800936 logging.warning('Received unsolicited response, ignored')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800937
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800938 def ParseMessage(self, buf, single=True):
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800939 if single:
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800940 try:
941 index = buf.index(_SEPARATOR)
942 except ValueError:
943 self._sock.UnRecv(buf)
944 return
945
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800946 msgs_json = [buf[:index]]
947 self._sock.UnRecv(buf[index + 2:])
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800948 else:
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800949 msgs_json = buf.split(_SEPARATOR)
950 self._sock.UnRecv(msgs_json.pop())
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800951
952 for msg_json in msgs_json:
953 try:
954 msg = json.loads(msg_json)
955 except ValueError:
956 # Ignore mal-formed message.
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800957 logging.error('mal-formed JSON request, ignored')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800958 continue
959
960 if 'name' in msg:
961 self.HandleRequest(msg)
962 elif 'response' in msg:
963 self.HandleResponse(msg)
964 else: # Ingnore mal-formed message.
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800965 logging.error('mal-formed JSON request, ignored')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800966
967 def ScanForTimeoutRequests(self):
Joel Kitching22b89042015-08-06 18:23:29 +0800968 """Scans for pending requests which have timed out.
969
970 If any timed-out requests are discovered, their handler is called with the
971 special response value of None.
972 """
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800973 for rid in self._requests.keys()[:]:
974 request_time, timeout, handler = self._requests[rid]
975 if self.Timestamp() - request_time > timeout:
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800976 if callable(handler):
977 handler(None)
978 else:
979 logging.error('Request %s timeout', rid)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800980 del self._requests[rid]
981
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800982 def InitiateDownload(self):
983 ttyname, filename = self._download_queue.get()
Wei-Ning Huangd521f282015-08-07 05:28:04 +0800984 sid = self._ttyname_to_sid[ttyname]
985 self.SpawnGhost(self.FILE, terminal_sid=sid,
Wei-Ning Huangae923642015-09-24 14:08:09 +0800986 file_op=('download', filename))
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800987
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800988 def Listen(self):
989 try:
990 while True:
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800991 rds, unused_wd, unused_xd = select.select([self._sock], [], [],
992 _PING_INTERVAL / 2)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800993
994 if self._sock in rds:
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800995 data = self._sock.Recv(_BUFSIZE)
Wei-Ning Huang09c19612015-11-24 16:29:09 +0800996
997 # Socket is closed
998 if len(data) == 0:
Wei-Ning Huang09c19612015-11-24 16:29:09 +0800999 break
1000
Wei-Ning Huanga28cd232016-01-27 15:04:41 +08001001 self.ParseMessage(data, self._register_status != SUCCESS)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001002
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001003 if (self._mode == self.AGENT and
1004 self.Timestamp() - self._last_ping > _PING_INTERVAL):
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001005 self.Ping()
1006 self.ScanForTimeoutRequests()
1007
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001008 if not self._download_queue.empty():
1009 self.InitiateDownload()
1010
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001011 if self._reset.is_set():
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001012 break
1013 except socket.error:
1014 raise RuntimeError('Connection dropped')
1015 except PingTimeoutError:
1016 raise RuntimeError('Connection timeout')
1017 finally:
Wei-Ning Huangf5311a02016-02-04 15:23:46 +08001018 self.Reset()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001019
1020 self._queue.put('resume')
1021
1022 if self._mode != Ghost.AGENT:
1023 sys.exit(1)
1024
1025 def Register(self):
1026 non_local = {}
1027 for addr in self._overlord_addrs:
1028 non_local['addr'] = addr
1029 def registered(response):
1030 if response is None:
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001031 self._reset.set()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001032 raise RuntimeError('Register request timeout')
Wei-Ning Huang63c16092015-09-18 16:20:27 +08001033
Wei-Ning Huang7ec55342015-09-17 08:46:06 +08001034 self._register_status = response['response']
1035 if response['response'] != SUCCESS:
1036 self._reset.set()
Peter Shih220a96d2016-12-22 17:02:16 +08001037 raise RuntimeError('Register: ' + response['response'])
Wei-Ning Huang7ec55342015-09-17 08:46:06 +08001038 else:
1039 logging.info('Registered with Overlord at %s:%d', *non_local['addr'])
1040 self._connected_addr = non_local['addr']
1041 self.Upgrade() # Check for upgrade
1042 self._queue.put('pause', True)
Wei-Ning Huang63c16092015-09-18 16:20:27 +08001043
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001044 try:
1045 logging.info('Trying %s:%d ...', *addr)
1046 self.Reset()
Wei-Ning Huangf5311a02016-02-04 15:23:46 +08001047
Peter Shih220a96d2016-12-22 17:02:16 +08001048 # Check if server has TLS enabled. Only check if self._tls_mode is
1049 # None.
Wei-Ning Huangb6605d22016-06-22 17:33:37 +08001050 # Only control channel needs to determine if TLS is enabled. Other mode
1051 # should use the TLSSettings passed in when it was spawned.
1052 if self._mode == Ghost.AGENT:
Peter Shih220a96d2016-12-22 17:02:16 +08001053 self._tls_settings.SetEnabled(
1054 self.TLSEnabled(*addr) if self._tls_mode is None
1055 else self._tls_mode)
Wei-Ning Huang47c79b82016-05-24 01:24:46 +08001056
Wei-Ning Huangf5311a02016-02-04 15:23:46 +08001057 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
1058 sock.settimeout(_CONNECT_TIMEOUT)
1059
1060 try:
1061 if self._tls_settings.Enabled():
1062 tls_context = self._tls_settings.Context()
1063 sock = tls_context.wrap_socket(sock, server_hostname=addr[0])
1064
1065 sock.connect(addr)
1066 except (ssl.SSLError, ssl.CertificateError) as e:
1067 logging.error('%s: %s', e.__class__.__name__, e)
1068 continue
1069 except IOError as e:
1070 if e.errno == 2: # No such file or directory
1071 logging.error('%s: %s', e.__class__.__name__, e)
1072 continue
1073 raise
1074
1075 self._sock = BufferedSocket(sock)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001076
1077 logging.info('Connection established, registering...')
1078 handler = {
1079 Ghost.AGENT: registered,
Wei-Ning Huangb8461202015-09-01 20:07:41 +08001080 Ghost.TERMINAL: self.SpawnTTYServer,
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001081 Ghost.SHELL: self.SpawnShellServer,
1082 Ghost.FILE: self.InitiateFileOperation,
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +08001083 Ghost.FORWARD: self.SpawnPortForwardServer,
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001084 }[self._mode]
1085
1086 # Machine ID may change if MAC address is used (USB-ethernet dongle
1087 # plugged/unplugged)
1088 self._machine_id = self.GetMachineID()
Wei-Ning Huang7d029b12015-03-06 10:32:15 +08001089 self.SendRequest('register',
1090 {'mode': self._mode, 'mid': self._machine_id,
Wei-Ning Huangfed95862015-08-07 03:17:11 +08001091 'sid': self._session_id,
Wei-Ning Huang7d029b12015-03-06 10:32:15 +08001092 'properties': self._properties}, handler)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001093 except socket.error:
1094 pass
1095 else:
Wei-Ning Huangf5311a02016-02-04 15:23:46 +08001096 sock.settimeout(None)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001097 self.Listen()
1098
Moja Hsuc9ecc8b2015-07-13 11:39:17 +08001099 raise RuntimeError('Cannot connect to any server')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001100
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001101 def Reconnect(self):
1102 logging.info('Received reconnect request from RPC server, reconnecting...')
1103 self._reset.set()
1104
Wei-Ning Huang7ec55342015-09-17 08:46:06 +08001105 def GetStatus(self):
1106 return self._register_status
1107
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001108 def AddToDownloadQueue(self, ttyname, filename):
1109 self._download_queue.put((ttyname, filename))
1110
Wei-Ning Huangd521f282015-08-07 05:28:04 +08001111 def RegisterTTY(self, session_id, ttyname):
1112 self._ttyname_to_sid[ttyname] = session_id
Wei-Ning Huange2981862015-08-03 15:03:08 +08001113
1114 def RegisterSession(self, session_id, process_id):
1115 self._terminal_sid_to_pid[session_id] = process_id
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001116
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001117 def StartLanDiscovery(self):
1118 """Start to listen to LAN discovery packet at
1119 _OVERLORD_LAN_DISCOVERY_PORT."""
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001120
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001121 def thread_func():
1122 s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
1123 s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
1124 s.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001125 try:
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001126 s.bind(('0.0.0.0', _OVERLORD_LAN_DISCOVERY_PORT))
1127 except socket.error as e:
Moja Hsuc9ecc8b2015-07-13 11:39:17 +08001128 logging.error('LAN discovery: %s, abort', e)
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001129 return
1130
1131 logging.info('LAN Discovery: started')
1132 while True:
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +08001133 rd, unused_wd, unused_xd = select.select([s], [], [], 1)
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001134
1135 if s in rd:
1136 data, source_addr = s.recvfrom(_BUFSIZE)
1137 parts = data.split()
1138 if parts[0] == 'OVERLORD':
1139 ip, port = parts[1].split(':')
1140 if not ip:
1141 ip = source_addr[0]
1142 self._queue.put((ip, int(port)), True)
1143
1144 try:
1145 obj = self._queue.get(False)
1146 except Queue.Empty:
1147 pass
1148 else:
1149 if type(obj) is not str:
1150 self._queue.put(obj)
1151 elif obj == 'pause':
1152 logging.info('LAN Discovery: paused')
1153 while obj != 'resume':
1154 obj = self._queue.get(True)
1155 logging.info('LAN Discovery: resumed')
1156
1157 t = threading.Thread(target=thread_func)
1158 t.daemon = True
1159 t.start()
1160
1161 def StartRPCServer(self):
Joel Kitching22b89042015-08-06 18:23:29 +08001162 logging.info('RPC Server: started')
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001163 rpc_server = SimpleJSONRPCServer((_DEFAULT_BIND_ADDRESS, _GHOST_RPC_PORT),
1164 logRequests=False)
1165 rpc_server.register_function(self.Reconnect, 'Reconnect')
Wei-Ning Huang7ec55342015-09-17 08:46:06 +08001166 rpc_server.register_function(self.GetStatus, 'GetStatus')
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001167 rpc_server.register_function(self.RegisterTTY, 'RegisterTTY')
Wei-Ning Huange2981862015-08-03 15:03:08 +08001168 rpc_server.register_function(self.RegisterSession, 'RegisterSession')
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001169 rpc_server.register_function(self.AddToDownloadQueue, 'AddToDownloadQueue')
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001170 t = threading.Thread(target=rpc_server.serve_forever)
1171 t.daemon = True
1172 t.start()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001173
Wei-Ning Huang829e0c82015-05-26 14:37:23 +08001174 def ScanServer(self):
1175 for meth in [self.GetGateWayIP, self.GetShopfloorIP]:
1176 for addr in [(x, _OVERLORD_PORT) for x in meth()]:
1177 if addr not in self._overlord_addrs:
1178 self._overlord_addrs.append(addr)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001179
Wei-Ning Huang11c35022015-10-21 16:52:32 +08001180 def Start(self, lan_disc=False, rpc_server=False):
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001181 logging.info('%s started', self.MODE_NAME[self._mode])
1182 logging.info('MID: %s', self._machine_id)
Wei-Ning Huangfed95862015-08-07 03:17:11 +08001183 logging.info('SID: %s', self._session_id)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001184
Wei-Ning Huangb05cde32015-08-01 09:48:41 +08001185 # We don't care about child process's return code, not wait is needed. This
1186 # is used to prevent zombie process from lingering in the system.
1187 self.SetIgnoreChild(True)
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001188
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001189 if lan_disc:
1190 self.StartLanDiscovery()
1191
1192 if rpc_server:
1193 self.StartRPCServer()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001194
1195 try:
1196 while True:
1197 try:
1198 addr = self._queue.get(False)
1199 except Queue.Empty:
1200 pass
1201 else:
1202 if type(addr) == tuple and addr not in self._overlord_addrs:
1203 logging.info('LAN Discovery: got overlord address %s:%d', *addr)
1204 self._overlord_addrs.append(addr)
1205
1206 try:
Wei-Ning Huang829e0c82015-05-26 14:37:23 +08001207 self.ScanServer()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001208 self.Register()
Joel Kitching22b89042015-08-06 18:23:29 +08001209 # Don't show stack trace for RuntimeError, which we use in this file for
1210 # plausible and expected errors (such as can't connect to server).
1211 except RuntimeError as e:
Wei-Ning Huang7ec55342015-09-17 08:46:06 +08001212 logging.info('%s, retrying in %ds', e.message, _RETRY_INTERVAL)
Joel Kitching22b89042015-08-06 18:23:29 +08001213 time.sleep(_RETRY_INTERVAL)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001214 except Exception as e:
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +08001215 unused_x, unused_y, exc_traceback = sys.exc_info()
Joel Kitching22b89042015-08-06 18:23:29 +08001216 traceback.print_tb(exc_traceback)
1217 logging.info('%s: %s, retrying in %ds',
1218 e.__class__.__name__, e.message, _RETRY_INTERVAL)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001219 time.sleep(_RETRY_INTERVAL)
1220
1221 self.Reset()
1222 except KeyboardInterrupt:
1223 logging.error('Received keyboard interrupt, quit')
1224 sys.exit(0)
1225
1226
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001227def GhostRPCServer():
Wei-Ning Huang8037c182015-09-19 04:41:50 +08001228 """Returns handler to Ghost's JSON RPC server."""
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001229 return jsonrpclib.Server('http://localhost:%d' % _GHOST_RPC_PORT)
1230
1231
Wei-Ning Huang8037c182015-09-19 04:41:50 +08001232def ForkToBackground():
1233 """Fork process to run in background."""
1234 pid = os.fork()
1235 if pid != 0:
1236 logging.info('Ghost(%d) running in background.', pid)
1237 sys.exit(0)
1238
1239
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001240def DownloadFile(filename):
Wei-Ning Huang8037c182015-09-19 04:41:50 +08001241 """Initiate a client-initiated file download."""
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001242 filepath = os.path.abspath(filename)
1243 if not os.path.exists(filepath):
Joel Kitching22b89042015-08-06 18:23:29 +08001244 logging.error('file `%s\' does not exist', filename)
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001245 sys.exit(1)
1246
1247 # Check if we actually have permission to read the file
1248 if not os.access(filepath, os.R_OK):
Joel Kitching22b89042015-08-06 18:23:29 +08001249 logging.error('can not open %s for reading', filepath)
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001250 sys.exit(1)
1251
1252 server = GhostRPCServer()
1253 server.AddToDownloadQueue(os.ttyname(0), filepath)
1254 sys.exit(0)
1255
1256
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001257def main():
Wei-Ning Huang5f3fa8f2015-10-24 15:08:48 +08001258 # Setup logging format
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001259 logger = logging.getLogger()
1260 logger.setLevel(logging.INFO)
Wei-Ning Huang5f3fa8f2015-10-24 15:08:48 +08001261 handler = logging.StreamHandler()
1262 formatter = logging.Formatter('%(asctime)s %(message)s', '%Y/%m/%d %H:%M:%S')
1263 handler.setFormatter(formatter)
1264 logger.addHandler(handler)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001265
Wei-Ning Huang7d029b12015-03-06 10:32:15 +08001266 parser = argparse.ArgumentParser()
Wei-Ning Huang8037c182015-09-19 04:41:50 +08001267 parser.add_argument('--fork', dest='fork', action='store_true', default=False,
1268 help='fork procecess to run in background')
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +08001269 parser.add_argument('--mid', metavar='MID', dest='mid', action='store',
1270 default=None, help='use MID as machine ID')
1271 parser.add_argument('--rand-mid', dest='mid', action='store_const',
1272 const=Ghost.RANDOM_MID, help='use random machine ID')
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001273 parser.add_argument('--no-lan-disc', dest='lan_disc', action='store_false',
1274 default=True, help='disable LAN discovery')
1275 parser.add_argument('--no-rpc-server', dest='rpc_server',
1276 action='store_false', default=True,
1277 help='disable RPC server')
Peter Shih220a96d2016-12-22 17:02:16 +08001278 parser.add_argument('--tls', dest='tls_mode', default='detect',
1279 choices=('y', 'n', 'detect'),
1280 help="specify 'y' or 'n' to force enable/disable TLS")
Wei-Ning Huangf5311a02016-02-04 15:23:46 +08001281 parser.add_argument('--tls-cert-file', metavar='TLS_CERT_FILE',
1282 dest='tls_cert_file', type=str, default=None,
1283 help='file containing the server TLS certificate in PEM '
1284 'format')
Wei-Ning Huang47c79b82016-05-24 01:24:46 +08001285 parser.add_argument('--tls-no-verify', dest='tls_no_verify',
1286 action='store_true', default=False,
1287 help='do not verify certificate if TLS is enabled')
Joel Kitching22b89042015-08-06 18:23:29 +08001288 parser.add_argument('--prop-file', metavar='PROP_FILE', dest='prop_file',
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001289 type=str, default=None,
Wei-Ning Huang7d029b12015-03-06 10:32:15 +08001290 help='file containing the JSON representation of client '
1291 'properties')
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001292 parser.add_argument('--download', metavar='FILE', dest='download', type=str,
1293 default=None, help='file to download')
Wei-Ning Huang23ed0162015-09-18 14:42:03 +08001294 parser.add_argument('--reset', dest='reset', default=False,
1295 action='store_true',
1296 help='reset ghost and reload all configs')
Wei-Ning Huang7d029b12015-03-06 10:32:15 +08001297 parser.add_argument('overlord_ip', metavar='OVERLORD_IP', type=str,
1298 nargs='*', help='overlord server address')
1299 args = parser.parse_args()
1300
Wei-Ning Huang8037c182015-09-19 04:41:50 +08001301 if args.fork:
1302 ForkToBackground()
1303
Wei-Ning Huang23ed0162015-09-18 14:42:03 +08001304 if args.reset:
1305 GhostRPCServer().Reconnect()
1306 sys.exit()
1307
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001308 if args.download:
1309 DownloadFile(args.download)
1310
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001311 addrs = [('localhost', _OVERLORD_PORT)]
Wei-Ning Huang47c79b82016-05-24 01:24:46 +08001312 addrs = [(x, _OVERLORD_PORT) for x in args.overlord_ip] + addrs
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001313
Wei-Ning Huangf5311a02016-02-04 15:23:46 +08001314 prop_file = os.path.abspath(args.prop_file) if args.prop_file else None
1315
Wei-Ning Huang47c79b82016-05-24 01:24:46 +08001316 tls_settings = TLSSettings(args.tls_cert_file, not args.tls_no_verify)
Peter Shih220a96d2016-12-22 17:02:16 +08001317 tls_mode = args.tls_mode
1318 tls_mode = {'y': True, 'n': False, 'detect': None}[tls_mode]
Wei-Ning Huangf5311a02016-02-04 15:23:46 +08001319 g = Ghost(addrs, tls_settings, Ghost.AGENT, args.mid,
Peter Shih220a96d2016-12-22 17:02:16 +08001320 prop_file=prop_file, tls_mode=tls_mode)
Wei-Ning Huang11c35022015-10-21 16:52:32 +08001321 g.Start(args.lan_disc, args.rpc_server)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001322
1323
1324if __name__ == '__main__':
Wei-Ning Huangf5311a02016-02-04 15:23:46 +08001325 try:
1326 main()
1327 except Exception as e:
1328 logging.error(e)