blob: c62c8fbc1bdb5348fa24f895e5c9314bdf389a3c [file] [log] [blame]
Mike Frysinger63bb3c72019-09-01 15:16:26 -04001#!/usr/bin/env python2
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08002# Copyright 2015 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
Peter Shih5cafebb2017-06-30 16:36:22 +08006from __future__ import print_function
7
Wei-Ning Huang7d029b12015-03-06 10:32:15 +08008import argparse
Yilin Yangf9fe1932019-11-04 17:09:34 +08009import binascii
10import codecs
Wei-Ning Huangb05cde32015-08-01 09:48:41 +080011import contextlib
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +080012import ctypes
13import ctypes.util
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080014import fcntl
Wei-Ning Huangb05cde32015-08-01 09:48:41 +080015import hashlib
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080016import json
17import logging
18import os
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +080019import platform
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080020import Queue
Wei-Ning Huang829e0c82015-05-26 14:37:23 +080021import re
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080022import select
Wei-Ning Huanga301f572015-06-03 17:34:21 +080023import signal
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080024import socket
Wei-Ning Huangf5311a02016-02-04 15:23:46 +080025import ssl
Wei-Ning Huangb05cde32015-08-01 09:48:41 +080026import struct
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080027import subprocess
28import sys
Moja Hsuc9ecc8b2015-07-13 11:39:17 +080029import termios
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080030import threading
31import time
Joel Kitching22b89042015-08-06 18:23:29 +080032import traceback
Wei-Ning Huang39169902015-09-19 06:00:23 +080033import tty
Wei-Ning Huange0def6a2015-11-05 15:41:24 +080034import urllib2
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080035import uuid
36
Wei-Ning Huang2132de32015-04-13 17:24:38 +080037import jsonrpclib
38from jsonrpclib.SimpleJSONRPCServer import SimpleJSONRPCServer
Yilin Yangf9fe1932019-11-04 17:09:34 +080039from six import PY2
Wei-Ning Huang2132de32015-04-13 17:24:38 +080040
Peter Shihc56c5b62016-12-22 12:30:57 +080041_GHOST_RPC_PORT = int(os.getenv('GHOST_RPC_PORT', 4499))
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080042
Peter Shihc56c5b62016-12-22 12:30:57 +080043_OVERLORD_PORT = int(os.getenv('OVERLORD_PORT', 4455))
44_OVERLORD_LAN_DISCOVERY_PORT = int(os.getenv('OVERLORD_LD_PORT', 4456))
45_OVERLORD_HTTP_PORT = int(os.getenv('OVERLORD_HTTP_PORT', 9000))
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080046
47_BUFSIZE = 8192
48_RETRY_INTERVAL = 2
49_SEPARATOR = '\r\n'
50_PING_TIMEOUT = 3
51_PING_INTERVAL = 5
52_REQUEST_TIMEOUT_SECS = 60
53_SHELL = os.getenv('SHELL', '/bin/bash')
Wei-Ning Huang7dbf4a72016-03-02 20:16:20 +080054_DEFAULT_BIND_ADDRESS = 'localhost'
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080055
Moja Hsuc9ecc8b2015-07-13 11:39:17 +080056_CONTROL_START = 128
57_CONTROL_END = 129
58
Wei-Ning Huanga301f572015-06-03 17:34:21 +080059_BLOCK_SIZE = 4096
Wei-Ning Huange0def6a2015-11-05 15:41:24 +080060_CONNECT_TIMEOUT = 3
Wei-Ning Huanga301f572015-06-03 17:34:21 +080061
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +080062# Stream control
63_STDIN_CLOSED = '##STDIN_CLOSED##'
64
Wei-Ning Huang7ec55342015-09-17 08:46:06 +080065SUCCESS = 'success'
66FAILED = 'failed'
67DISCONNECTED = 'disconnected'
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080068
Joel Kitching22b89042015-08-06 18:23:29 +080069
Wei-Ning Huang1cea6112015-03-02 12:45:34 +080070class PingTimeoutError(Exception):
71 pass
72
73
74class RequestError(Exception):
75 pass
76
77
Wei-Ning Huangf5311a02016-02-04 15:23:46 +080078class BufferedSocket(object):
Wei-Ning Huanga28cd232016-01-27 15:04:41 +080079 """A buffered socket that supports unrecv.
80
81 Allow putting back data back to the socket for the next recv() call.
82 """
Wei-Ning Huangf5311a02016-02-04 15:23:46 +080083 def __init__(self, sock):
84 self.sock = sock
Wei-Ning Huanga28cd232016-01-27 15:04:41 +080085 self._buf = ''
86
Wei-Ning Huangf5311a02016-02-04 15:23:46 +080087 def fileno(self):
88 return self.sock.fileno()
89
Wei-Ning Huanga28cd232016-01-27 15:04:41 +080090 def Recv(self, bufsize, flags=0):
91 if self._buf:
92 if len(self._buf) >= bufsize:
93 ret = self._buf[:bufsize]
94 self._buf = self._buf[bufsize:]
95 return ret
96 else:
97 ret = self._buf
98 self._buf = ''
Wei-Ning Huangf5311a02016-02-04 15:23:46 +080099 return ret + self.sock.recv(bufsize - len(ret), flags)
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800100 else:
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800101 return self.sock.recv(bufsize, flags)
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800102
103 def UnRecv(self, buf):
104 self._buf = buf + self._buf
105
106 def Send(self, *args, **kwargs):
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800107 return self.sock.send(*args, **kwargs)
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800108
109 def RecvBuf(self):
110 """Only recive from buffer."""
111 ret = self._buf
112 self._buf = ''
113 return ret
114
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800115 def Close(self):
116 self.sock.close()
117
118
119class TLSSettings(object):
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800120 def __init__(self, tls_cert_file, verify):
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800121 """Constructor.
122
123 Args:
124 tls_cert_file: TLS certificate in PEM format.
125 enable_tls_without_verify: enable TLS but don't verify certificate.
126 """
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800127 self._enabled = False
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800128 self._tls_cert_file = tls_cert_file
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800129 self._verify = verify
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800130 self._tls_context = None
131
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800132 def _UpdateContext(self):
133 if not self._enabled:
134 self._tls_context = None
135 return
136
137 self._tls_context = ssl.SSLContext(ssl.PROTOCOL_TLSv1_2)
138 self._tls_context.verify_mode = ssl.CERT_REQUIRED
139
140 if self._verify:
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800141 if self._tls_cert_file:
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800142 self._tls_context.check_hostname = True
143 try:
144 self._tls_context.load_verify_locations(self._tls_cert_file)
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800145 logging.info('TLSSettings: using user-supplied ca-certificate')
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800146 except IOError as e:
147 logging.error('TLSSettings: %s: %s', self._tls_cert_file, e)
148 sys.exit(1)
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800149 else:
150 self._tls_context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
151 logging.info('TLSSettings: using built-in ca-certificates')
152 else:
153 self._tls_context.verify_mode = ssl.CERT_NONE
154 logging.info('TLSSettings: skipping TLS verification!!!')
155
156 def SetEnabled(self, enabled):
157 logging.info('TLSSettings: enabled: %s', enabled)
158
159 if self._enabled != enabled:
160 self._enabled = enabled
161 self._UpdateContext()
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800162
163 def Enabled(self):
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800164 return self._enabled
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800165
166 def Context(self):
167 return self._tls_context
168
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800169
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800170class Ghost(object):
171 """Ghost implements the client protocol of Overlord.
172
173 Ghost provide terminal/shell/logcat functionality and manages the client
174 side connectivity.
175 """
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800176 NONE, AGENT, TERMINAL, SHELL, LOGCAT, FILE, FORWARD = range(7)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800177
178 MODE_NAME = {
179 NONE: 'NONE',
180 AGENT: 'Agent',
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800181 TERMINAL: 'Terminal',
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800182 SHELL: 'Shell',
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800183 LOGCAT: 'Logcat',
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800184 FILE: 'File',
185 FORWARD: 'Forward'
Peter Shihe6afab32018-09-11 17:16:48 +0800186 }
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800187
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +0800188 RANDOM_MID = '##random_mid##'
189
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800190 def __init__(self, overlord_addrs, tls_settings=None, mode=AGENT, mid=None,
191 sid=None, prop_file=None, terminal_sid=None, tty_device=None,
Peter Shih220a96d2016-12-22 17:02:16 +0800192 command=None, file_op=None, port=None, tls_mode=None):
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800193 """Constructor.
194
195 Args:
196 overlord_addrs: a list of possible address of overlord.
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800197 tls_settings: a TLSSetting object.
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800198 mode: client mode, either AGENT, SHELL or LOGCAT
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +0800199 mid: a str to set for machine ID. If mid equals Ghost.RANDOM_MID, machine
200 id is randomly generated.
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800201 sid: session ID. If the connection is requested by overlord, sid should
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800202 be set to the corresponding session id assigned by overlord.
Wei-Ning Huang23ed0162015-09-18 14:42:03 +0800203 prop_file: properties file filename.
Wei-Ning Huangd521f282015-08-07 05:28:04 +0800204 terminal_sid: the terminal session ID associate with this client. This is
205 use for file download.
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800206 tty_device: the terminal device to open, if tty_device is None, as pseudo
207 terminal will be opened instead.
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800208 command: the command to execute when we are in SHELL mode.
Wei-Ning Huang8ee3bcd2015-10-01 17:10:01 +0800209 file_op: a tuple (action, filepath, perm). action is either 'download' or
210 'upload'. perm is the permission to set for the file.
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800211 port: port number to forward.
Peter Shih220a96d2016-12-22 17:02:16 +0800212 tls_mode: can be [True, False, None]. if not None, skip detection of
213 TLS and assume whether server use TLS or not.
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800214 """
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800215 assert mode in [Ghost.AGENT, Ghost.TERMINAL, Ghost.SHELL, Ghost.FILE,
216 Ghost.FORWARD]
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800217 if mode == Ghost.SHELL:
218 assert command is not None
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800219 if mode == Ghost.FILE:
220 assert file_op is not None
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800221
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800222 self._platform = platform.system()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800223 self._overlord_addrs = overlord_addrs
Wei-Ning Huangad330c52015-03-12 20:34:18 +0800224 self._connected_addr = None
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800225 self._tls_settings = tls_settings
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +0800226 self._mid = mid
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800227 self._sock = None
Wei-Ning Huang7ec55342015-09-17 08:46:06 +0800228 self._mode = mode
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800229 self._machine_id = self.GetMachineID()
Wei-Ning Huangfed95862015-08-07 03:17:11 +0800230 self._session_id = sid if sid is not None else str(uuid.uuid4())
Wei-Ning Huangd521f282015-08-07 05:28:04 +0800231 self._terminal_session_id = terminal_sid
Wei-Ning Huang23ed0162015-09-18 14:42:03 +0800232 self._ttyname_to_sid = {}
233 self._terminal_sid_to_pid = {}
234 self._prop_file = prop_file
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800235 self._properties = {}
Wei-Ning Huang7ec55342015-09-17 08:46:06 +0800236 self._register_status = DISCONNECTED
237 self._reset = threading.Event()
Peter Shih220a96d2016-12-22 17:02:16 +0800238 self._tls_mode = tls_mode
Wei-Ning Huang7ec55342015-09-17 08:46:06 +0800239
240 # RPC
Wei-Ning Huang7ec55342015-09-17 08:46:06 +0800241 self._requests = {}
242 self._queue = Queue.Queue()
243
244 # Protocol specific
245 self._last_ping = 0
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800246 self._tty_device = tty_device
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800247 self._shell_command = command
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800248 self._file_op = file_op
Wei-Ning Huang7ec55342015-09-17 08:46:06 +0800249 self._download_queue = Queue.Queue()
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800250 self._port = port
Wei-Ning Huang7ec55342015-09-17 08:46:06 +0800251
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800252 def SetIgnoreChild(self, status):
253 # Only ignore child for Agent since only it could spawn child Ghost.
254 if self._mode == Ghost.AGENT:
255 signal.signal(signal.SIGCHLD,
256 signal.SIG_IGN if status else signal.SIG_DFL)
257
258 def GetFileSha1(self, filename):
259 with open(filename, 'r') as f:
260 return hashlib.sha1(f.read()).hexdigest()
261
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800262 def TLSEnabled(self, host, port):
263 """Determine if TLS is enabled on given server address."""
Wei-Ning Huang58833882015-09-16 16:52:37 +0800264 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
265 try:
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800266 # Allow any certificate since we only want to check if server talks TLS.
267 context = ssl.SSLContext(ssl.PROTOCOL_TLSv1_2)
268 context.verify_mode = ssl.CERT_NONE
Wei-Ning Huang58833882015-09-16 16:52:37 +0800269
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800270 sock = context.wrap_socket(sock, server_hostname=host)
271 sock.settimeout(_CONNECT_TIMEOUT)
272 sock.connect((host, port))
273 return True
Wei-Ning Huangb6605d22016-06-22 17:33:37 +0800274 except ssl.SSLError:
275 return False
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800276 except socket.error: # Connect refused or timeout
277 raise
Wei-Ning Huang58833882015-09-16 16:52:37 +0800278 except Exception:
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800279 return False # For whatever reason above failed, assume False
Wei-Ning Huang58833882015-09-16 16:52:37 +0800280
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800281 def Upgrade(self):
282 logging.info('Upgrade: initiating upgrade sequence...')
283
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800284 try:
Wei-Ning Huangb6605d22016-06-22 17:33:37 +0800285 https_enabled = self.TLSEnabled(self._connected_addr[0],
286 _OVERLORD_HTTP_PORT)
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800287 except socket.error:
Wei-Ning Huangb6605d22016-06-22 17:33:37 +0800288 logging.error('Upgrade: failed to connect to Overlord HTTP server, '
289 'abort')
290 return
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800291
292 if self._tls_settings.Enabled() and not https_enabled:
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800293 logging.error('Upgrade: TLS enforced but found Overlord HTTP server '
294 'without TLS enabled! Possible mis-configuration or '
295 'DNS/IP spoofing detected, abort')
296 return
297
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800298 scriptpath = os.path.abspath(sys.argv[0])
Wei-Ning Huang03f9f762015-09-16 21:51:35 +0800299 url = 'http%s://%s:%d/upgrade/ghost.py' % (
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800300 's' if https_enabled else '', self._connected_addr[0],
Wei-Ning Huang03f9f762015-09-16 21:51:35 +0800301 _OVERLORD_HTTP_PORT)
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800302
303 # Download sha1sum for ghost.py for verification
304 try:
Wei-Ning Huange0def6a2015-11-05 15:41:24 +0800305 with contextlib.closing(
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800306 urllib2.urlopen(url + '.sha1', timeout=_CONNECT_TIMEOUT,
307 context=self._tls_settings.Context())) as f:
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800308 if f.getcode() != 200:
309 raise RuntimeError('HTTP status %d' % f.getcode())
310 sha1sum = f.read().strip()
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800311 except (ssl.SSLError, ssl.CertificateError) as e:
312 logging.error('Upgrade: %s: %s', e.__class__.__name__, e)
313 return
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800314 except Exception:
315 logging.error('Upgrade: failed to download sha1sum file, abort')
316 return
317
318 if self.GetFileSha1(scriptpath) == sha1sum:
319 logging.info('Upgrade: ghost is already up-to-date, skipping upgrade')
320 return
321
322 # Download upgrade version of ghost.py
323 try:
Wei-Ning Huange0def6a2015-11-05 15:41:24 +0800324 with contextlib.closing(
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800325 urllib2.urlopen(url, timeout=_CONNECT_TIMEOUT,
326 context=self._tls_settings.Context())) as f:
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800327 if f.getcode() != 200:
328 raise RuntimeError('HTTP status %d' % f.getcode())
329 data = f.read()
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800330 except (ssl.SSLError, ssl.CertificateError) as e:
331 logging.error('Upgrade: %s: %s', e.__class__.__name__, e)
332 return
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800333 except Exception:
334 logging.error('Upgrade: failed to download upgrade, abort')
335 return
336
337 # Compare SHA1 sum
338 if hashlib.sha1(data).hexdigest() != sha1sum:
339 logging.error('Upgrade: sha1sum mismatch, abort')
340 return
341
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800342 try:
343 with open(scriptpath, 'w') as f:
344 f.write(data)
345 except Exception:
346 logging.error('Upgrade: failed to write upgrade onto disk, abort')
347 return
348
349 logging.info('Upgrade: restarting ghost...')
350 self.CloseSockets()
351 self.SetIgnoreChild(False)
Wei-Ning Huang5f3fa8f2015-10-24 15:08:48 +0800352 os.execve(scriptpath, [scriptpath] + sys.argv[1:], os.environ)
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800353
Wei-Ning Huang23ed0162015-09-18 14:42:03 +0800354 def LoadProperties(self):
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800355 try:
Wei-Ning Huang23ed0162015-09-18 14:42:03 +0800356 if self._prop_file:
357 with open(self._prop_file, 'r') as f:
358 self._properties = json.loads(f.read())
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800359 except Exception as e:
Peter Shih769b0772018-02-26 14:44:28 +0800360 logging.error('LoadProperties: %s', e)
Wei-Ning Huang7d029b12015-03-06 10:32:15 +0800361
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800362 def CloseSockets(self):
363 # Close sockets opened by parent process, since we don't use it anymore.
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800364 if self._platform == 'Linux':
365 for fd in os.listdir('/proc/self/fd/'):
366 try:
367 real_fd = os.readlink('/proc/self/fd/%s' % fd)
368 if real_fd.startswith('socket'):
369 os.close(int(fd))
370 except Exception:
371 pass
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800372
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800373 def SpawnGhost(self, mode, sid=None, terminal_sid=None, tty_device=None,
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800374 command=None, file_op=None, port=None):
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800375 """Spawn a child ghost with specific mode.
376
377 Returns:
378 The spawned child process pid.
379 """
Joel Kitching22b89042015-08-06 18:23:29 +0800380 # Restore the default signal handler, so our child won't have problems.
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800381 self.SetIgnoreChild(False)
382
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800383 pid = os.fork()
384 if pid == 0:
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800385 self.CloseSockets()
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800386 g = Ghost([self._connected_addr], tls_settings=self._tls_settings,
387 mode=mode, mid=Ghost.RANDOM_MID, sid=sid,
Wei-Ning Huang23ed0162015-09-18 14:42:03 +0800388 terminal_sid=terminal_sid, tty_device=tty_device,
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800389 command=command, file_op=file_op, port=port)
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800390 g.Start()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800391 sys.exit(0)
392 else:
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800393 self.SetIgnoreChild(True)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800394 return pid
395
396 def Timestamp(self):
397 return int(time.time())
398
399 def GetGateWayIP(self):
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800400 if self._platform == 'Darwin':
401 output = subprocess.check_output(['route', '-n', 'get', 'default'])
402 ret = re.search('gateway: (.*)', output)
403 if ret:
404 return [ret.group(1)]
Peter Shiha78867d2018-02-26 14:17:51 +0800405 return []
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800406 elif self._platform == 'Linux':
407 with open('/proc/net/route', 'r') as f:
408 lines = f.readlines()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800409
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800410 ips = []
411 for line in lines:
412 parts = line.split('\t')
413 if parts[2] == '00000000':
414 continue
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800415
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800416 try:
Yilin Yangf9fe1932019-11-04 17:09:34 +0800417 h = codecs.decode(parts[2], 'hex')
418 # TODO(kerker) Remove when py3 upgrade complete
419 if PY2:
420 ips.append('%d.%d.%d.%d' % tuple(ord(x) for x in reversed(h)))
421 else:
422 ips.append('.'.join([str(x) for x in reversed(h)]))
423 except (TypeError, binascii.Error):
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800424 pass
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800425
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800426 return ips
427 else:
428 logging.warning('GetGateWayIP: unsupported platform')
429 return []
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800430
Hung-Te Lin41ff8f32017-08-30 08:10:39 +0800431 def GetFactoryServerIP(self):
Wei-Ning Huang829e0c82015-05-26 14:37:23 +0800432 try:
Peter Shihcb0e5512017-06-14 16:59:46 +0800433 import factory_common # pylint: disable=unused-variable
Hung-Te Lin41ff8f32017-08-30 08:10:39 +0800434 from cros.factory.test import server_proxy
Wei-Ning Huang829e0c82015-05-26 14:37:23 +0800435
Hung-Te Lin41ff8f32017-08-30 08:10:39 +0800436 url = server_proxy.GetServerURL()
Wei-Ning Huang829e0c82015-05-26 14:37:23 +0800437 match = re.match(r'^https?://(.*):.*$', url)
438 if match:
439 return [match.group(1)]
440 except Exception:
441 pass
442 return []
443
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800444 def GetMachineID(self):
445 """Generates machine-dependent ID string for a machine.
446 There are many ways to generate a machine ID:
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800447 Linux:
448 1. factory device_id
Peter Shih5f1f48c2017-06-26 14:12:00 +0800449 2. /sys/class/dmi/id/product_uuid (only available on intel machines)
450 3. MAC address
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800451 We follow the listed order to generate machine ID, and fallback to the
452 next alternative if the previous doesn't work.
453
454 Darwin:
455 All Darwin system should have the IOPlatformSerialNumber attribute.
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800456 """
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +0800457 if self._mid == Ghost.RANDOM_MID:
Wei-Ning Huangaed90452015-03-23 17:50:21 +0800458 return str(uuid.uuid4())
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +0800459 elif self._mid:
460 return self._mid
Wei-Ning Huangaed90452015-03-23 17:50:21 +0800461
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800462 # Darwin
463 if self._platform == 'Darwin':
464 output = subprocess.check_output(['ioreg', '-rd1', '-c',
465 'IOPlatformExpertDevice'])
466 ret = re.search('"IOPlatformSerialNumber" = "(.*)"', output)
467 if ret:
468 return ret.group(1)
469
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800470 # Try factory device id
471 try:
Peter Shihcb0e5512017-06-14 16:59:46 +0800472 import factory_common # pylint: disable=unused-variable
Hung-Te Linda8eb992017-09-28 03:27:12 +0800473 from cros.factory.test import session
474 return session.GetDeviceID()
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800475 except Exception:
476 pass
477
Wei-Ning Huang1d7603b2015-07-03 17:38:56 +0800478 # Try DMI product UUID
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800479 try:
480 with open('/sys/class/dmi/id/product_uuid', 'r') as f:
481 return f.read().strip()
482 except Exception:
483 pass
484
Wei-Ning Huang1d7603b2015-07-03 17:38:56 +0800485 # Use MAC address if non is available
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800486 try:
487 macs = []
488 ifaces = sorted(os.listdir('/sys/class/net'))
489 for iface in ifaces:
490 if iface == 'lo':
491 continue
492
493 with open('/sys/class/net/%s/address' % iface, 'r') as f:
494 macs.append(f.read().strip())
495
496 return ';'.join(macs)
497 except Exception:
498 pass
499
Peter Shihcb0e5512017-06-14 16:59:46 +0800500 raise RuntimeError("can't generate machine ID")
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800501
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800502 def GetProcessWorkingDirectory(self, pid):
503 if self._platform == 'Linux':
504 return os.readlink('/proc/%d/cwd' % pid)
505 elif self._platform == 'Darwin':
506 PROC_PIDVNODEPATHINFO = 9
507 proc_vnodepathinfo_size = 2352
508 vid_path_offset = 152
509
510 proc = ctypes.cdll.LoadLibrary(ctypes.util.find_library('libproc'))
511 buf = ctypes.create_string_buffer('\0' * proc_vnodepathinfo_size)
512 proc.proc_pidinfo(pid, PROC_PIDVNODEPATHINFO, 0,
513 ctypes.byref(buf), proc_vnodepathinfo_size)
514 buf = buf.raw[vid_path_offset:]
515 n = buf.index('\0')
516 return buf[:n]
517 else:
518 raise RuntimeError('GetProcessWorkingDirectory: unsupported platform')
519
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800520 def Reset(self):
521 """Reset state and clear request handlers."""
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800522 if self._sock is not None:
523 self._sock.Close()
524 self._sock = None
Wei-Ning Huang2132de32015-04-13 17:24:38 +0800525 self._reset.clear()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800526 self._last_ping = 0
527 self._requests = {}
Wei-Ning Huang23ed0162015-09-18 14:42:03 +0800528 self.LoadProperties()
Wei-Ning Huang7ec55342015-09-17 08:46:06 +0800529 self._register_status = DISCONNECTED
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800530
531 def SendMessage(self, msg):
532 """Serialize the message and send it through the socket."""
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800533 self._sock.Send(json.dumps(msg) + _SEPARATOR)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800534
535 def SendRequest(self, name, args, handler=None,
536 timeout=_REQUEST_TIMEOUT_SECS):
537 if handler and not callable(handler):
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800538 raise RequestError('Invalid request handler for msg "%s"' % name)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800539
540 rid = str(uuid.uuid4())
541 msg = {'rid': rid, 'timeout': timeout, 'name': name, 'params': args}
Wei-Ning Huange2981862015-08-03 15:03:08 +0800542 if timeout >= 0:
543 self._requests[rid] = [self.Timestamp(), timeout, handler]
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800544 self.SendMessage(msg)
545
546 def SendResponse(self, omsg, status, params=None):
547 msg = {'rid': omsg['rid'], 'response': status, 'params': params}
548 self.SendMessage(msg)
549
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800550 def HandleTTYControl(self, fd, control_str):
551 msg = json.loads(control_str)
Moja Hsuc9ecc8b2015-07-13 11:39:17 +0800552 command = msg['command']
553 params = msg['params']
554 if command == 'resize':
555 # some error happened on websocket
556 if len(params) != 2:
557 return
558 winsize = struct.pack('HHHH', params[0], params[1], 0, 0)
559 fcntl.ioctl(fd, termios.TIOCSWINSZ, winsize)
560 else:
561 logging.warn('Invalid request command "%s"', command)
562
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800563 def SpawnTTYServer(self, unused_var):
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800564 """Spawn a TTY server and forward I/O to the TCP socket."""
565 logging.info('SpawnTTYServer: started')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800566
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800567 try:
568 if self._tty_device is None:
569 pid, fd = os.forkpty()
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800570
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800571 if pid == 0:
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800572 ttyname = os.ttyname(sys.stdout.fileno())
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800573 try:
574 server = GhostRPCServer()
575 server.RegisterTTY(self._session_id, ttyname)
576 server.RegisterSession(self._session_id, os.getpid())
577 except Exception:
578 # If ghost is launched without RPC server, the call will fail but we
579 # can ignore it.
580 pass
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800581
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800582 # The directory that contains the current running ghost script
583 script_dir = os.path.dirname(os.path.abspath(sys.argv[0]))
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800584
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800585 env = os.environ.copy()
586 env['USER'] = os.getenv('USER', 'root')
587 env['HOME'] = os.getenv('HOME', '/root')
588 env['PATH'] = os.getenv('PATH') + ':%s' % script_dir
589 os.chdir(env['HOME'])
590 os.execve(_SHELL, [_SHELL], env)
591 else:
592 fd = os.open(self._tty_device, os.O_RDWR)
Wei-Ning Huang39169902015-09-19 06:00:23 +0800593 tty.setraw(fd)
594 attr = termios.tcgetattr(fd)
595 attr[0] &= ~(termios.IXON | termios.IXOFF)
596 attr[2] |= termios.CLOCAL
597 attr[2] &= ~termios.CRTSCTS
598 attr[4] = termios.B115200
599 attr[5] = termios.B115200
600 termios.tcsetattr(fd, termios.TCSANOW, attr)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800601
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800602 nonlocals = {'control_state': None, 'control_str': ''}
603
604 def _ProcessBuffer(buf):
605 write_buffer = ''
606 while buf:
607 if nonlocals['control_state']:
608 if chr(_CONTROL_END) in buf:
609 index = buf.index(chr(_CONTROL_END))
610 nonlocals['control_str'] += buf[:index]
611 self.HandleTTYControl(fd, nonlocals['control_str'])
612 nonlocals['control_state'] = None
613 nonlocals['control_str'] = ''
614 buf = buf[index+1:]
615 else:
616 nonlocals['control_str'] += buf
617 buf = ''
618 else:
619 if chr(_CONTROL_START) in buf:
620 nonlocals['control_state'] = _CONTROL_START
621 index = buf.index(chr(_CONTROL_START))
622 write_buffer += buf[:index]
623 buf = buf[index+1:]
624 else:
625 write_buffer += buf
626 buf = ''
627
628 if write_buffer:
629 os.write(fd, write_buffer)
630
631 _ProcessBuffer(self._sock.RecvBuf())
632
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800633 while True:
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800634 rd, unused_wd, unused_xd = select.select([self._sock, fd], [], [])
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800635
636 if fd in rd:
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800637 self._sock.Send(os.read(fd, _BUFSIZE))
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800638
639 if self._sock in rd:
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800640 buf = self._sock.Recv(_BUFSIZE)
Peter Shihaacbc2f2017-06-16 14:39:29 +0800641 if not buf:
Wei-Ning Huang5f3fa8f2015-10-24 15:08:48 +0800642 raise RuntimeError('connection terminated')
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800643 _ProcessBuffer(buf)
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800644 except Exception as e:
645 logging.error('SpawnTTYServer: %s', e)
646 finally:
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800647 self._sock.Close()
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800648
649 logging.info('SpawnTTYServer: terminated')
650 sys.exit(0)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800651
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800652 def SpawnShellServer(self, unused_var):
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800653 """Spawn a shell server and forward input/output from/to the TCP socket."""
654 logging.info('SpawnShellServer: started')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800655
Wei-Ning Huang5f3fa8f2015-10-24 15:08:48 +0800656 # Add ghost executable to PATH
657 script_dir = os.path.dirname(os.path.abspath(sys.argv[0]))
658 env = os.environ.copy()
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800659 env['PATH'] = '%s:%s' % (script_dir, os.getenv('PATH'))
Wei-Ning Huang5f3fa8f2015-10-24 15:08:48 +0800660
661 # Execute shell command from HOME directory
662 os.chdir(os.getenv('HOME', '/tmp'))
663
Wei-Ning Huang0f4a5372015-03-09 15:12:07 +0800664 p = subprocess.Popen(self._shell_command, stdin=subprocess.PIPE,
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800665 stdout=subprocess.PIPE, stderr=subprocess.PIPE,
Wei-Ning Huang5f3fa8f2015-10-24 15:08:48 +0800666 shell=True, env=env)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800667
668 def make_non_block(fd):
669 fl = fcntl.fcntl(fd, fcntl.F_GETFL)
670 fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
671
672 make_non_block(p.stdout)
673 make_non_block(p.stderr)
674
675 try:
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800676 p.stdin.write(self._sock.RecvBuf())
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800677
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800678 while True:
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800679 rd, unused_wd, unused_xd = select.select(
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800680 [p.stdout, p.stderr, self._sock], [], [])
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800681 if p.stdout in rd:
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800682 self._sock.Send(p.stdout.read(_BUFSIZE))
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800683
684 if p.stderr in rd:
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800685 self._sock.Send(p.stderr.read(_BUFSIZE))
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800686
687 if self._sock in rd:
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800688 ret = self._sock.Recv(_BUFSIZE)
Peter Shihaacbc2f2017-06-16 14:39:29 +0800689 if not ret:
Wei-Ning Huang5f3fa8f2015-10-24 15:08:48 +0800690 raise RuntimeError('connection terminated')
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800691
692 try:
693 idx = ret.index(_STDIN_CLOSED * 2)
694 p.stdin.write(ret[:idx])
695 p.stdin.close()
696 except ValueError:
697 p.stdin.write(ret)
Wei-Ning Huangf14c84e2015-08-03 15:03:08 +0800698 p.poll()
Peter Shihe6afab32018-09-11 17:16:48 +0800699 if p.returncode is not None:
Wei-Ning Huangf14c84e2015-08-03 15:03:08 +0800700 break
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800701 except Exception as e:
702 logging.error('SpawnShellServer: %s', e)
Wei-Ning Huangf14c84e2015-08-03 15:03:08 +0800703 finally:
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800704 # Check if the process is terminated. If not, Send SIGTERM to process,
705 # then wait for 1 second. Send another SIGKILL to make sure the process is
706 # terminated.
707 p.poll()
708 if p.returncode is None:
709 try:
710 p.terminate()
711 time.sleep(1)
712 p.kill()
713 except Exception:
714 pass
Wei-Ning Huang5f3fa8f2015-10-24 15:08:48 +0800715
716 p.wait()
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800717 self._sock.Close()
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800718
719 logging.info('SpawnShellServer: terminated')
720 sys.exit(0)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800721
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800722 def InitiateFileOperation(self, unused_var):
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800723 if self._file_op[0] == 'download':
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800724 try:
725 size = os.stat(self._file_op[1]).st_size
726 except OSError as e:
727 logging.error('InitiateFileOperation: download: %s', e)
728 sys.exit(1)
729
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800730 self.SendRequest('request_to_download',
Wei-Ning Huangd521f282015-08-07 05:28:04 +0800731 {'terminal_sid': self._terminal_session_id,
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800732 'filename': os.path.basename(self._file_op[1]),
733 'size': size})
Wei-Ning Huange2981862015-08-03 15:03:08 +0800734 elif self._file_op[0] == 'upload':
735 self.SendRequest('clear_to_upload', {}, timeout=-1)
736 self.StartUploadServer()
737 else:
738 logging.error('InitiateFileOperation: unknown file operation, ignored')
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800739
740 def StartDownloadServer(self):
741 logging.info('StartDownloadServer: started')
742
743 try:
744 with open(self._file_op[1], 'rb') as f:
745 while True:
746 data = f.read(_BLOCK_SIZE)
Peter Shihaacbc2f2017-06-16 14:39:29 +0800747 if not data:
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800748 break
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800749 self._sock.Send(data)
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800750 except Exception as e:
751 logging.error('StartDownloadServer: %s', e)
752 finally:
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800753 self._sock.Close()
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800754
755 logging.info('StartDownloadServer: terminated')
756 sys.exit(0)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800757
Wei-Ning Huange2981862015-08-03 15:03:08 +0800758 def StartUploadServer(self):
759 logging.info('StartUploadServer: started')
Wei-Ning Huange2981862015-08-03 15:03:08 +0800760 try:
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800761 filepath = self._file_op[1]
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800762 dirname = os.path.dirname(filepath)
763 if not os.path.exists(dirname):
764 try:
765 os.makedirs(dirname)
766 except Exception:
767 pass
Wei-Ning Huange2981862015-08-03 15:03:08 +0800768
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800769 with open(filepath, 'wb') as f:
Wei-Ning Huang8ee3bcd2015-10-01 17:10:01 +0800770 if self._file_op[2]:
771 os.fchmod(f.fileno(), self._file_op[2])
772
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800773 f.write(self._sock.RecvBuf())
774
Wei-Ning Huange2981862015-08-03 15:03:08 +0800775 while True:
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800776 rd, unused_wd, unused_xd = select.select([self._sock], [], [])
Wei-Ning Huange2981862015-08-03 15:03:08 +0800777 if self._sock in rd:
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800778 buf = self._sock.Recv(_BLOCK_SIZE)
Peter Shihaacbc2f2017-06-16 14:39:29 +0800779 if not buf:
Wei-Ning Huange2981862015-08-03 15:03:08 +0800780 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:
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800787 self._sock.Close()
Wei-Ning Huange2981862015-08-03 15:03:08 +0800788
789 logging.info('StartUploadServer: terminated')
790 sys.exit(0)
791
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800792 def SpawnPortForwardServer(self, unused_var):
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800793 """Spawn a port forwarding server and forward I/O to the TCP socket."""
794 logging.info('SpawnPortForwardServer: started')
795
796 src_sock = None
797 try:
798 src_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
Wei-Ning Huange0def6a2015-11-05 15:41:24 +0800799 src_sock.settimeout(_CONNECT_TIMEOUT)
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800800 src_sock.connect(('localhost', self._port))
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800801
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800802 src_sock.send(self._sock.RecvBuf())
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800803
804 while True:
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800805 rd, unused_wd, unused_xd = select.select([self._sock, src_sock], [], [])
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800806
807 if self._sock in rd:
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800808 data = self._sock.Recv(_BUFSIZE)
Peter Shihaacbc2f2017-06-16 14:39:29 +0800809 if not data:
Wei-Ning Huang5f3fa8f2015-10-24 15:08:48 +0800810 raise RuntimeError('connection terminated')
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800811 src_sock.send(data)
812
813 if src_sock in rd:
814 data = src_sock.recv(_BUFSIZE)
Peter Shihaacbc2f2017-06-16 14:39:29 +0800815 if not data:
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800816 break
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800817 self._sock.Send(data)
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800818 except Exception as e:
819 logging.error('SpawnPortForwardServer: %s', e)
820 finally:
821 if src_sock:
822 src_sock.close()
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800823 self._sock.Close()
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800824
825 logging.info('SpawnPortForwardServer: terminated')
826 sys.exit(0)
827
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800828 def Ping(self):
829 def timeout_handler(x):
830 if x is None:
831 raise PingTimeoutError
832
833 self._last_ping = self.Timestamp()
834 self.SendRequest('ping', {}, timeout_handler, 5)
835
Wei-Ning Huangae923642015-09-24 14:08:09 +0800836 def HandleFileDownloadRequest(self, msg):
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800837 params = msg['params']
Wei-Ning Huangae923642015-09-24 14:08:09 +0800838 filepath = params['filename']
839 if not os.path.isabs(filepath):
840 filepath = os.path.join(os.getenv('HOME', '/tmp'), filepath)
841
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800842 try:
Wei-Ning Huang11c35022015-10-21 16:52:32 +0800843 with open(filepath, 'r') as _:
Wei-Ning Huang46a3fc92015-10-06 02:35:27 +0800844 pass
845 except Exception as e:
Peter Shiha78867d2018-02-26 14:17:51 +0800846 self.SendResponse(msg, str(e))
847 return
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800848
849 self.SpawnGhost(self.FILE, params['sid'],
Wei-Ning Huangae923642015-09-24 14:08:09 +0800850 file_op=('download', filepath))
851 self.SendResponse(msg, SUCCESS)
852
853 def HandleFileUploadRequest(self, msg):
854 params = msg['params']
855
856 # Resolve upload filepath
857 filename = params['filename']
858 dest_path = filename
859
860 # If dest is specified, use it first
861 dest_path = params.get('dest', '')
862 if dest_path:
863 if not os.path.isabs(dest_path):
864 dest_path = os.path.join(os.getenv('HOME', '/tmp'), dest_path)
865
866 if os.path.isdir(dest_path):
867 dest_path = os.path.join(dest_path, filename)
868 else:
869 target_dir = os.getenv('HOME', '/tmp')
870
871 # Terminal session ID found, upload to it's current working directory
Peter Shihe6afab32018-09-11 17:16:48 +0800872 if 'terminal_sid' in params:
Wei-Ning Huangae923642015-09-24 14:08:09 +0800873 pid = self._terminal_sid_to_pid.get(params['terminal_sid'], None)
874 if pid:
Wei-Ning Huanga0e55b82016-02-10 14:32:07 +0800875 try:
876 target_dir = self.GetProcessWorkingDirectory(pid)
877 except Exception as e:
878 logging.error(e)
Wei-Ning Huangae923642015-09-24 14:08:09 +0800879
880 dest_path = os.path.join(target_dir, filename)
881
882 try:
883 os.makedirs(os.path.dirname(dest_path))
884 except Exception:
885 pass
886
887 try:
888 with open(dest_path, 'w') as _:
889 pass
890 except Exception as e:
Peter Shiha78867d2018-02-26 14:17:51 +0800891 self.SendResponse(msg, str(e))
892 return
Wei-Ning Huangae923642015-09-24 14:08:09 +0800893
Wei-Ning Huangd6f69762015-10-01 21:02:07 +0800894 # If not check_only, spawn FILE mode ghost agent to handle upload
895 if not params.get('check_only', False):
896 self.SpawnGhost(self.FILE, params['sid'],
897 file_op=('upload', dest_path, params.get('perm', None)))
Wei-Ning Huang7ec55342015-09-17 08:46:06 +0800898 self.SendResponse(msg, SUCCESS)
Wei-Ning Huang552cd702015-08-12 16:11:13 +0800899
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800900 def HandleRequest(self, msg):
Wei-Ning Huange2981862015-08-03 15:03:08 +0800901 command = msg['name']
902 params = msg['params']
903
904 if command == 'upgrade':
Wei-Ning Huangb05cde32015-08-01 09:48:41 +0800905 self.Upgrade()
Wei-Ning Huange2981862015-08-03 15:03:08 +0800906 elif command == 'terminal':
Wei-Ning Huangb8461202015-09-01 20:07:41 +0800907 self.SpawnGhost(self.TERMINAL, params['sid'],
908 tty_device=params['tty_device'])
Wei-Ning Huang7ec55342015-09-17 08:46:06 +0800909 self.SendResponse(msg, SUCCESS)
Wei-Ning Huange2981862015-08-03 15:03:08 +0800910 elif command == 'shell':
911 self.SpawnGhost(self.SHELL, params['sid'], command=params['command'])
Wei-Ning Huang7ec55342015-09-17 08:46:06 +0800912 self.SendResponse(msg, SUCCESS)
Wei-Ning Huange2981862015-08-03 15:03:08 +0800913 elif command == 'file_download':
Wei-Ning Huangae923642015-09-24 14:08:09 +0800914 self.HandleFileDownloadRequest(msg)
Wei-Ning Huange2981862015-08-03 15:03:08 +0800915 elif command == 'clear_to_download':
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800916 self.StartDownloadServer()
Wei-Ning Huange2981862015-08-03 15:03:08 +0800917 elif command == 'file_upload':
Wei-Ning Huangae923642015-09-24 14:08:09 +0800918 self.HandleFileUploadRequest(msg)
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800919 elif command == 'forward':
920 self.SpawnGhost(self.FORWARD, params['sid'], port=params['port'])
921 self.SendResponse(msg, SUCCESS)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800922
923 def HandleResponse(self, response):
924 rid = str(response['rid'])
925 if rid in self._requests:
926 handler = self._requests[rid][2]
927 del self._requests[rid]
928 if callable(handler):
929 handler(response)
930 else:
Joel Kitching22b89042015-08-06 18:23:29 +0800931 logging.warning('Received unsolicited response, ignored')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800932
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800933 def ParseMessage(self, buf, single=True):
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800934 if single:
Wei-Ning Huangf5311a02016-02-04 15:23:46 +0800935 try:
936 index = buf.index(_SEPARATOR)
937 except ValueError:
938 self._sock.UnRecv(buf)
939 return
940
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800941 msgs_json = [buf[:index]]
942 self._sock.UnRecv(buf[index + 2:])
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800943 else:
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800944 msgs_json = buf.split(_SEPARATOR)
945 self._sock.UnRecv(msgs_json.pop())
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800946
947 for msg_json in msgs_json:
948 try:
949 msg = json.loads(msg_json)
950 except ValueError:
951 # Ignore mal-formed message.
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800952 logging.error('mal-formed JSON request, ignored')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800953 continue
954
955 if 'name' in msg:
956 self.HandleRequest(msg)
957 elif 'response' in msg:
958 self.HandleResponse(msg)
959 else: # Ingnore mal-formed message.
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +0800960 logging.error('mal-formed JSON request, ignored')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800961
962 def ScanForTimeoutRequests(self):
Joel Kitching22b89042015-08-06 18:23:29 +0800963 """Scans for pending requests which have timed out.
964
965 If any timed-out requests are discovered, their handler is called with the
966 special response value of None.
967 """
Yilin Yang78fa12e2019-09-25 14:21:10 +0800968 for rid in list(self._requests):
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800969 request_time, timeout, handler = self._requests[rid]
970 if self.Timestamp() - request_time > timeout:
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800971 if callable(handler):
972 handler(None)
973 else:
974 logging.error('Request %s timeout', rid)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800975 del self._requests[rid]
976
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800977 def InitiateDownload(self):
978 ttyname, filename = self._download_queue.get()
Wei-Ning Huangd521f282015-08-07 05:28:04 +0800979 sid = self._ttyname_to_sid[ttyname]
980 self.SpawnGhost(self.FILE, terminal_sid=sid,
Wei-Ning Huangae923642015-09-24 14:08:09 +0800981 file_op=('download', filename))
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800982
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800983 def Listen(self):
984 try:
985 while True:
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800986 rds, unused_wd, unused_xd = select.select([self._sock], [], [],
Yilin Yang14d02a22019-11-01 11:32:03 +0800987 _PING_INTERVAL // 2)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800988
989 if self._sock in rds:
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800990 data = self._sock.Recv(_BUFSIZE)
Wei-Ning Huang09c19612015-11-24 16:29:09 +0800991
992 # Socket is closed
Peter Shihaacbc2f2017-06-16 14:39:29 +0800993 if not data:
Wei-Ning Huang09c19612015-11-24 16:29:09 +0800994 break
995
Wei-Ning Huanga28cd232016-01-27 15:04:41 +0800996 self.ParseMessage(data, self._register_status != SUCCESS)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +0800997
Wei-Ning Huanga301f572015-06-03 17:34:21 +0800998 if (self._mode == self.AGENT and
999 self.Timestamp() - self._last_ping > _PING_INTERVAL):
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001000 self.Ping()
1001 self.ScanForTimeoutRequests()
1002
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001003 if not self._download_queue.empty():
1004 self.InitiateDownload()
1005
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001006 if self._reset.is_set():
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001007 break
1008 except socket.error:
1009 raise RuntimeError('Connection dropped')
1010 except PingTimeoutError:
1011 raise RuntimeError('Connection timeout')
1012 finally:
Wei-Ning Huangf5311a02016-02-04 15:23:46 +08001013 self.Reset()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001014
1015 self._queue.put('resume')
1016
1017 if self._mode != Ghost.AGENT:
1018 sys.exit(1)
1019
1020 def Register(self):
1021 non_local = {}
1022 for addr in self._overlord_addrs:
1023 non_local['addr'] = addr
1024 def registered(response):
1025 if response is None:
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001026 self._reset.set()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001027 raise RuntimeError('Register request timeout')
Wei-Ning Huang63c16092015-09-18 16:20:27 +08001028
Wei-Ning Huang7ec55342015-09-17 08:46:06 +08001029 self._register_status = response['response']
1030 if response['response'] != SUCCESS:
1031 self._reset.set()
Peter Shih220a96d2016-12-22 17:02:16 +08001032 raise RuntimeError('Register: ' + response['response'])
Wei-Ning Huang7ec55342015-09-17 08:46:06 +08001033 else:
1034 logging.info('Registered with Overlord at %s:%d', *non_local['addr'])
1035 self._connected_addr = non_local['addr']
1036 self.Upgrade() # Check for upgrade
1037 self._queue.put('pause', True)
Wei-Ning Huang63c16092015-09-18 16:20:27 +08001038
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001039 try:
1040 logging.info('Trying %s:%d ...', *addr)
1041 self.Reset()
Wei-Ning Huangf5311a02016-02-04 15:23:46 +08001042
Peter Shih220a96d2016-12-22 17:02:16 +08001043 # Check if server has TLS enabled. Only check if self._tls_mode is
1044 # None.
Wei-Ning Huangb6605d22016-06-22 17:33:37 +08001045 # Only control channel needs to determine if TLS is enabled. Other mode
1046 # should use the TLSSettings passed in when it was spawned.
1047 if self._mode == Ghost.AGENT:
Peter Shih220a96d2016-12-22 17:02:16 +08001048 self._tls_settings.SetEnabled(
1049 self.TLSEnabled(*addr) if self._tls_mode is None
1050 else self._tls_mode)
Wei-Ning Huang47c79b82016-05-24 01:24:46 +08001051
Wei-Ning Huangf5311a02016-02-04 15:23:46 +08001052 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
1053 sock.settimeout(_CONNECT_TIMEOUT)
1054
1055 try:
1056 if self._tls_settings.Enabled():
1057 tls_context = self._tls_settings.Context()
1058 sock = tls_context.wrap_socket(sock, server_hostname=addr[0])
1059
1060 sock.connect(addr)
1061 except (ssl.SSLError, ssl.CertificateError) as e:
1062 logging.error('%s: %s', e.__class__.__name__, e)
1063 continue
1064 except IOError as e:
1065 if e.errno == 2: # No such file or directory
1066 logging.error('%s: %s', e.__class__.__name__, e)
1067 continue
1068 raise
1069
1070 self._sock = BufferedSocket(sock)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001071
1072 logging.info('Connection established, registering...')
1073 handler = {
1074 Ghost.AGENT: registered,
Wei-Ning Huangb8461202015-09-01 20:07:41 +08001075 Ghost.TERMINAL: self.SpawnTTYServer,
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001076 Ghost.SHELL: self.SpawnShellServer,
1077 Ghost.FILE: self.InitiateFileOperation,
Wei-Ning Huangdadbeb62015-09-20 00:38:27 +08001078 Ghost.FORWARD: self.SpawnPortForwardServer,
Peter Shihe6afab32018-09-11 17:16:48 +08001079 }[self._mode]
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001080
1081 # Machine ID may change if MAC address is used (USB-ethernet dongle
1082 # plugged/unplugged)
1083 self._machine_id = self.GetMachineID()
Wei-Ning Huang7d029b12015-03-06 10:32:15 +08001084 self.SendRequest('register',
1085 {'mode': self._mode, 'mid': self._machine_id,
Wei-Ning Huangfed95862015-08-07 03:17:11 +08001086 'sid': self._session_id,
Wei-Ning Huang7d029b12015-03-06 10:32:15 +08001087 'properties': self._properties}, handler)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001088 except socket.error:
1089 pass
1090 else:
Wei-Ning Huangf5311a02016-02-04 15:23:46 +08001091 sock.settimeout(None)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001092 self.Listen()
1093
Moja Hsuc9ecc8b2015-07-13 11:39:17 +08001094 raise RuntimeError('Cannot connect to any server')
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001095
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001096 def Reconnect(self):
1097 logging.info('Received reconnect request from RPC server, reconnecting...')
1098 self._reset.set()
1099
Wei-Ning Huang7ec55342015-09-17 08:46:06 +08001100 def GetStatus(self):
Peter Shih5cafebb2017-06-30 16:36:22 +08001101 status = self._register_status
1102 if self._register_status == SUCCESS:
1103 ip, port = self._sock.sock.getpeername()
1104 status += ' %s:%d' % (ip, port)
1105 return status
Wei-Ning Huang7ec55342015-09-17 08:46:06 +08001106
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001107 def AddToDownloadQueue(self, ttyname, filename):
1108 self._download_queue.put((ttyname, filename))
1109
Wei-Ning Huangd521f282015-08-07 05:28:04 +08001110 def RegisterTTY(self, session_id, ttyname):
1111 self._ttyname_to_sid[ttyname] = session_id
Wei-Ning Huange2981862015-08-03 15:03:08 +08001112
1113 def RegisterSession(self, session_id, process_id):
1114 self._terminal_sid_to_pid[session_id] = process_id
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001115
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001116 def StartLanDiscovery(self):
1117 """Start to listen to LAN discovery packet at
1118 _OVERLORD_LAN_DISCOVERY_PORT."""
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001119
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001120 def thread_func():
1121 s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
1122 s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
1123 s.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001124 try:
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001125 s.bind(('0.0.0.0', _OVERLORD_LAN_DISCOVERY_PORT))
1126 except socket.error as e:
Moja Hsuc9ecc8b2015-07-13 11:39:17 +08001127 logging.error('LAN discovery: %s, abort', e)
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001128 return
1129
1130 logging.info('LAN Discovery: started')
1131 while True:
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +08001132 rd, unused_wd, unused_xd = select.select([s], [], [], 1)
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001133
1134 if s in rd:
1135 data, source_addr = s.recvfrom(_BUFSIZE)
1136 parts = data.split()
1137 if parts[0] == 'OVERLORD':
1138 ip, port = parts[1].split(':')
1139 if not ip:
1140 ip = source_addr[0]
1141 self._queue.put((ip, int(port)), True)
1142
1143 try:
1144 obj = self._queue.get(False)
1145 except Queue.Empty:
1146 pass
1147 else:
Peter Shihaacbc2f2017-06-16 14:39:29 +08001148 if not isinstance(obj, str):
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001149 self._queue.put(obj)
1150 elif obj == 'pause':
1151 logging.info('LAN Discovery: paused')
1152 while obj != 'resume':
1153 obj = self._queue.get(True)
1154 logging.info('LAN Discovery: resumed')
1155
1156 t = threading.Thread(target=thread_func)
1157 t.daemon = True
1158 t.start()
1159
1160 def StartRPCServer(self):
Joel Kitching22b89042015-08-06 18:23:29 +08001161 logging.info('RPC Server: started')
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001162 rpc_server = SimpleJSONRPCServer((_DEFAULT_BIND_ADDRESS, _GHOST_RPC_PORT),
1163 logRequests=False)
1164 rpc_server.register_function(self.Reconnect, 'Reconnect')
Wei-Ning Huang7ec55342015-09-17 08:46:06 +08001165 rpc_server.register_function(self.GetStatus, 'GetStatus')
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001166 rpc_server.register_function(self.RegisterTTY, 'RegisterTTY')
Wei-Ning Huange2981862015-08-03 15:03:08 +08001167 rpc_server.register_function(self.RegisterSession, 'RegisterSession')
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001168 rpc_server.register_function(self.AddToDownloadQueue, 'AddToDownloadQueue')
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001169 t = threading.Thread(target=rpc_server.serve_forever)
1170 t.daemon = True
1171 t.start()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001172
Wei-Ning Huang829e0c82015-05-26 14:37:23 +08001173 def ScanServer(self):
Hung-Te Lin41ff8f32017-08-30 08:10:39 +08001174 for meth in [self.GetGateWayIP, self.GetFactoryServerIP]:
Wei-Ning Huang829e0c82015-05-26 14:37:23 +08001175 for addr in [(x, _OVERLORD_PORT) for x in meth()]:
1176 if addr not in self._overlord_addrs:
1177 self._overlord_addrs.append(addr)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001178
Wei-Ning Huang11c35022015-10-21 16:52:32 +08001179 def Start(self, lan_disc=False, rpc_server=False):
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001180 logging.info('%s started', self.MODE_NAME[self._mode])
1181 logging.info('MID: %s', self._machine_id)
Wei-Ning Huangfed95862015-08-07 03:17:11 +08001182 logging.info('SID: %s', self._session_id)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001183
Wei-Ning Huangb05cde32015-08-01 09:48:41 +08001184 # We don't care about child process's return code, not wait is needed. This
1185 # is used to prevent zombie process from lingering in the system.
1186 self.SetIgnoreChild(True)
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001187
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001188 if lan_disc:
1189 self.StartLanDiscovery()
1190
1191 if rpc_server:
1192 self.StartRPCServer()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001193
1194 try:
1195 while True:
1196 try:
1197 addr = self._queue.get(False)
1198 except Queue.Empty:
1199 pass
1200 else:
Peter Shihaacbc2f2017-06-16 14:39:29 +08001201 if isinstance(addr, tuple) and addr not in self._overlord_addrs:
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001202 logging.info('LAN Discovery: got overlord address %s:%d', *addr)
1203 self._overlord_addrs.append(addr)
1204
1205 try:
Wei-Ning Huang829e0c82015-05-26 14:37:23 +08001206 self.ScanServer()
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001207 self.Register()
Joel Kitching22b89042015-08-06 18:23:29 +08001208 # Don't show stack trace for RuntimeError, which we use in this file for
1209 # plausible and expected errors (such as can't connect to server).
1210 except RuntimeError as e:
Yilin Yang58948af2019-10-30 18:28:55 +08001211 logging.info('%s, retrying in %ds', str(e), _RETRY_INTERVAL)
Joel Kitching22b89042015-08-06 18:23:29 +08001212 time.sleep(_RETRY_INTERVAL)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001213 except Exception as e:
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +08001214 unused_x, unused_y, exc_traceback = sys.exc_info()
Joel Kitching22b89042015-08-06 18:23:29 +08001215 traceback.print_tb(exc_traceback)
1216 logging.info('%s: %s, retrying in %ds',
Yilin Yang58948af2019-10-30 18:28:55 +08001217 e.__class__.__name__, str(e), _RETRY_INTERVAL)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001218 time.sleep(_RETRY_INTERVAL)
1219
1220 self.Reset()
1221 except KeyboardInterrupt:
1222 logging.error('Received keyboard interrupt, quit')
1223 sys.exit(0)
1224
1225
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001226def GhostRPCServer():
Wei-Ning Huang8037c182015-09-19 04:41:50 +08001227 """Returns handler to Ghost's JSON RPC server."""
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001228 return jsonrpclib.Server('http://localhost:%d' % _GHOST_RPC_PORT)
1229
1230
Wei-Ning Huang8037c182015-09-19 04:41:50 +08001231def ForkToBackground():
1232 """Fork process to run in background."""
1233 pid = os.fork()
1234 if pid != 0:
1235 logging.info('Ghost(%d) running in background.', pid)
1236 sys.exit(0)
1237
1238
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001239def DownloadFile(filename):
Wei-Ning Huang8037c182015-09-19 04:41:50 +08001240 """Initiate a client-initiated file download."""
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001241 filepath = os.path.abspath(filename)
1242 if not os.path.exists(filepath):
Joel Kitching22b89042015-08-06 18:23:29 +08001243 logging.error('file `%s\' does not exist', filename)
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001244 sys.exit(1)
1245
1246 # Check if we actually have permission to read the file
1247 if not os.access(filepath, os.R_OK):
Joel Kitching22b89042015-08-06 18:23:29 +08001248 logging.error('can not open %s for reading', filepath)
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001249 sys.exit(1)
1250
1251 server = GhostRPCServer()
1252 server.AddToDownloadQueue(os.ttyname(0), filepath)
1253 sys.exit(0)
1254
1255
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001256def main():
Wei-Ning Huang5f3fa8f2015-10-24 15:08:48 +08001257 # Setup logging format
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001258 logger = logging.getLogger()
1259 logger.setLevel(logging.INFO)
Wei-Ning Huang5f3fa8f2015-10-24 15:08:48 +08001260 handler = logging.StreamHandler()
1261 formatter = logging.Formatter('%(asctime)s %(message)s', '%Y/%m/%d %H:%M:%S')
1262 handler.setFormatter(formatter)
1263 logger.addHandler(handler)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001264
Wei-Ning Huang7d029b12015-03-06 10:32:15 +08001265 parser = argparse.ArgumentParser()
Wei-Ning Huang8037c182015-09-19 04:41:50 +08001266 parser.add_argument('--fork', dest='fork', action='store_true', default=False,
1267 help='fork procecess to run in background')
Wei-Ning Huangc9c97f02015-05-19 15:05:42 +08001268 parser.add_argument('--mid', metavar='MID', dest='mid', action='store',
1269 default=None, help='use MID as machine ID')
1270 parser.add_argument('--rand-mid', dest='mid', action='store_const',
1271 const=Ghost.RANDOM_MID, help='use random machine ID')
Wei-Ning Huang2132de32015-04-13 17:24:38 +08001272 parser.add_argument('--no-lan-disc', dest='lan_disc', action='store_false',
1273 default=True, help='disable LAN discovery')
1274 parser.add_argument('--no-rpc-server', dest='rpc_server',
1275 action='store_false', default=True,
1276 help='disable RPC server')
Peter Shih220a96d2016-12-22 17:02:16 +08001277 parser.add_argument('--tls', dest='tls_mode', default='detect',
1278 choices=('y', 'n', 'detect'),
1279 help="specify 'y' or 'n' to force enable/disable TLS")
Wei-Ning Huangf5311a02016-02-04 15:23:46 +08001280 parser.add_argument('--tls-cert-file', metavar='TLS_CERT_FILE',
1281 dest='tls_cert_file', type=str, default=None,
1282 help='file containing the server TLS certificate in PEM '
1283 'format')
Wei-Ning Huang47c79b82016-05-24 01:24:46 +08001284 parser.add_argument('--tls-no-verify', dest='tls_no_verify',
1285 action='store_true', default=False,
1286 help='do not verify certificate if TLS is enabled')
Joel Kitching22b89042015-08-06 18:23:29 +08001287 parser.add_argument('--prop-file', metavar='PROP_FILE', dest='prop_file',
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001288 type=str, default=None,
Wei-Ning Huang7d029b12015-03-06 10:32:15 +08001289 help='file containing the JSON representation of client '
1290 'properties')
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001291 parser.add_argument('--download', metavar='FILE', dest='download', type=str,
1292 default=None, help='file to download')
Wei-Ning Huang23ed0162015-09-18 14:42:03 +08001293 parser.add_argument('--reset', dest='reset', default=False,
1294 action='store_true',
1295 help='reset ghost and reload all configs')
Peter Shih5cafebb2017-06-30 16:36:22 +08001296 parser.add_argument('--status', dest='status', default=False,
1297 action='store_true',
1298 help='show status of the client')
Wei-Ning Huang7d029b12015-03-06 10:32:15 +08001299 parser.add_argument('overlord_ip', metavar='OVERLORD_IP', type=str,
1300 nargs='*', help='overlord server address')
1301 args = parser.parse_args()
1302
Peter Shih5cafebb2017-06-30 16:36:22 +08001303 if args.status:
1304 print(GhostRPCServer().GetStatus())
1305 sys.exit()
1306
Wei-Ning Huang8037c182015-09-19 04:41:50 +08001307 if args.fork:
1308 ForkToBackground()
1309
Wei-Ning Huang23ed0162015-09-18 14:42:03 +08001310 if args.reset:
1311 GhostRPCServer().Reconnect()
1312 sys.exit()
1313
Wei-Ning Huanga301f572015-06-03 17:34:21 +08001314 if args.download:
1315 DownloadFile(args.download)
1316
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001317 addrs = [('localhost', _OVERLORD_PORT)]
Wei-Ning Huang47c79b82016-05-24 01:24:46 +08001318 addrs = [(x, _OVERLORD_PORT) for x in args.overlord_ip] + addrs
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001319
Wei-Ning Huangf5311a02016-02-04 15:23:46 +08001320 prop_file = os.path.abspath(args.prop_file) if args.prop_file else None
1321
Wei-Ning Huang47c79b82016-05-24 01:24:46 +08001322 tls_settings = TLSSettings(args.tls_cert_file, not args.tls_no_verify)
Peter Shih220a96d2016-12-22 17:02:16 +08001323 tls_mode = args.tls_mode
1324 tls_mode = {'y': True, 'n': False, 'detect': None}[tls_mode]
Wei-Ning Huangf5311a02016-02-04 15:23:46 +08001325 g = Ghost(addrs, tls_settings, Ghost.AGENT, args.mid,
Peter Shih220a96d2016-12-22 17:02:16 +08001326 prop_file=prop_file, tls_mode=tls_mode)
Wei-Ning Huang11c35022015-10-21 16:52:32 +08001327 g.Start(args.lan_disc, args.rpc_server)
Wei-Ning Huang1cea6112015-03-02 12:45:34 +08001328
1329
1330if __name__ == '__main__':
Wei-Ning Huangf5311a02016-02-04 15:23:46 +08001331 try:
1332 main()
1333 except Exception as e:
1334 logging.error(e)