blob: bf02d48d8364650fb0ca105cf4bc2eb011d26bc5 [file] [log] [blame]
Yilin Yang19da6932019-12-10 13:39:28 +08001#!/usr/bin/env python3
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +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
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08006import argparse
7import ast
8import base64
9import fcntl
Peter Shih13e78c52018-01-23 12:57:07 +080010import functools
Wei-Ning Huangba768ab2016-02-07 14:38:06 +080011import getpass
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +080012import hashlib
Yilin Yang752db712019-09-27 15:42:38 +080013import http.client
Yilin Yang8d4f9d02019-11-28 17:12:11 +080014from io import StringIO
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +080015import json
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +080016import logging
17import os
18import re
19import select
20import signal
21import socket
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +080022import ssl
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +080023import struct
24import subprocess
25import sys
26import tempfile
27import termios
28import threading
29import time
30import tty
Peter Shih99b73ec2017-06-16 17:54:15 +080031import unicodedata # required by pyinstaller, pylint: disable=unused-import
Yilin Yangf54fb912020-01-08 11:42:38 +080032import urllib.error
33import urllib.parse
34import urllib.request
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +080035
Peter Shih99b73ec2017-06-16 17:54:15 +080036import jsonrpclib
Stimim Chena30447c2020-10-06 10:04:00 +080037from jsonrpclib import config
Peter Shih99b73ec2017-06-16 17:54:15 +080038from jsonrpclib.SimpleJSONRPCServer import SimpleJSONRPCServer
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +080039from ws4py.client import WebSocketBaseClient
Peter Shih99b73ec2017-06-16 17:54:15 +080040import yaml
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +080041
Yilin Yang83c8f442020-05-05 13:46:51 +080042from cros.factory.utils import process_utils
43
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +080044
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +080045_CERT_DIR = os.path.expanduser('~/.config/ovl')
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +080046
47_ESCAPE = '~'
48_BUFSIZ = 8192
49_OVERLORD_PORT = 4455
50_OVERLORD_HTTP_PORT = 9000
51_OVERLORD_CLIENT_DAEMON_PORT = 4488
52_OVERLORD_CLIENT_DAEMON_RPC_ADDR = ('127.0.0.1', _OVERLORD_CLIENT_DAEMON_PORT)
53
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +080054_CONNECT_TIMEOUT = 3
Wei-Ning Huangee7ca8d2015-12-12 05:48:02 +080055_DEFAULT_HTTP_TIMEOUT = 30
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +080056_LIST_CACHE_TIMEOUT = 2
57_DEFAULT_TERMINAL_WIDTH = 80
Wei-Ning Huangee7ca8d2015-12-12 05:48:02 +080058_RETRY_TIMES = 3
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +080059
60# echo -n overlord | md5sum
61_HTTP_BOUNDARY_MAGIC = '9246f080c855a69012707ab53489b921'
62
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +080063# Terminal resize control
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +080064_CONTROL_START = 128
65_CONTROL_END = 129
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +080066
67# Stream control
68_STDIN_CLOSED = '##STDIN_CLOSED##'
69
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +080070_SSH_CONTROL_SOCKET_PREFIX = os.path.join(tempfile.gettempdir(),
71 'ovl-ssh-control-')
72
Wei-Ning Huang47c79b82016-05-24 01:24:46 +080073_TLS_CERT_FAILED_WARNING = """
74@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
75@ WARNING: REMOTE HOST VERIFICATION HAS FAILED! @
76@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
77IT IS POSSIBLE THAT SOMEONE IS DOING SOMETHING NASTY!
78Someone could be eavesdropping on you right now (man-in-the-middle attack)!
79It is also possible that the server is using a self-signed certificate.
80The fingerprint for the TLS host certificate sent by the remote host is
81
82%s
83
84Do you want to trust this certificate and proceed? [Y/n] """
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +080085
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +080086_TLS_CERT_CHANGED_WARNING = """
87@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
88@ WARNING: REMOTE HOST IDENTIFICATION HAS CHANGED! @
89@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
90IT IS POSSIBLE THAT SOMEONE IS DOING SOMETHING NASTY!
91Someone could be eavesdropping on you right now (man-in-the-middle attack)!
Wei-Ning Huang47c79b82016-05-24 01:24:46 +080092It is also possible that the TLS host certificate has just been changed.
93The fingerprint for the TLS host certificate sent by the remote host is
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +080094
95%s
96
97Remove '%s' if you still want to proceed.
98SSL Certificate verification failed."""
99
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800100
101def GetVersionDigest():
102 """Return the sha1sum of the current executing script."""
Wei-Ning Huang31609662016-08-11 00:22:25 +0800103 # Check python script by default
104 filename = __file__
105
106 # If we are running from a frozen binary, we should calculate the checksum
107 # against that binary instead of the python script.
108 # See: https://pyinstaller.readthedocs.io/en/stable/runtime-information.html
109 if getattr(sys, 'frozen', False):
110 filename = sys.executable
111
Yilin Yang0412c272019-12-05 16:57:40 +0800112 with open(filename, 'rb') as f:
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800113 return hashlib.sha1(f.read()).hexdigest()
114
115
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800116def GetTLSCertPath(host):
117 return os.path.join(_CERT_DIR, '%s.cert' % host)
118
119
120def UrlOpen(state, url):
Yilin Yangf54fb912020-01-08 11:42:38 +0800121 """Wrapper for urllib.request.urlopen.
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800122
123 It selects correct HTTP scheme according to self._state.ssl, add HTTP
124 basic auth headers, and add specify correct SSL context.
125 """
126 url = MakeRequestUrl(state, url)
Yilin Yangf54fb912020-01-08 11:42:38 +0800127 request = urllib.request.Request(url)
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800128 if state.username is not None and state.password is not None:
129 request.add_header(*BasicAuthHeader(state.username, state.password))
Yilin Yangf54fb912020-01-08 11:42:38 +0800130 return urllib.request.urlopen(request, timeout=_DEFAULT_HTTP_TIMEOUT,
131 context=state.ssl_context)
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800132
133
134def GetTLSCertificateSHA1Fingerprint(cert_pem):
135 beg = cert_pem.index('\n')
136 end = cert_pem.rindex('\n', 0, len(cert_pem) - 2)
137 cert_pem = cert_pem[beg:end] # Remove BEGIN/END CERTIFICATE boundary
138 cert_der = base64.b64decode(cert_pem)
139 return hashlib.sha1(cert_der).hexdigest()
140
141
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800142def KillGraceful(pid, wait_secs=1):
143 """Kill a process gracefully by first sending SIGTERM, wait for some time,
144 then send SIGKILL to make sure it's killed."""
145 try:
146 os.kill(pid, signal.SIGTERM)
147 time.sleep(wait_secs)
148 os.kill(pid, signal.SIGKILL)
149 except OSError:
150 pass
151
152
Wei-Ning Huangee7ca8d2015-12-12 05:48:02 +0800153def AutoRetry(action_name, retries):
154 """Decorator for retry function call."""
155 def Wrap(func):
Peter Shih13e78c52018-01-23 12:57:07 +0800156 @functools.wraps(func)
Wei-Ning Huangee7ca8d2015-12-12 05:48:02 +0800157 def Loop(*args, **kwargs):
Wei-Ning Huang5564eea2016-01-19 14:36:45 +0800158 for unused_i in range(retries):
Wei-Ning Huangee7ca8d2015-12-12 05:48:02 +0800159 try:
160 func(*args, **kwargs)
161 except Exception as e:
Wei-Ning Huang5564eea2016-01-19 14:36:45 +0800162 print('error: %s: %s: retrying ...' % (args[0], e))
Wei-Ning Huangee7ca8d2015-12-12 05:48:02 +0800163 else:
164 break
165 else:
Wei-Ning Huang5564eea2016-01-19 14:36:45 +0800166 print('error: failed to %s %s' % (action_name, args[0]))
Wei-Ning Huangee7ca8d2015-12-12 05:48:02 +0800167 return Loop
168 return Wrap
169
170
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800171def BasicAuthHeader(user, password):
172 """Return HTTP basic auth header."""
Stimim Chena30447c2020-10-06 10:04:00 +0800173 credential = base64.b64encode(
174 b'%s:%s' % (user.encode('utf-8'), password.encode('utf-8')))
175 return ('Authorization', 'Basic %s' % credential.decode('utf-8'))
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800176
177
178def GetTerminalSize():
179 """Retrieve terminal window size."""
180 ws = struct.pack('HHHH', 0, 0, 0, 0)
181 ws = fcntl.ioctl(0, termios.TIOCGWINSZ, ws)
182 lines, columns, unused_x, unused_y = struct.unpack('HHHH', ws)
183 return lines, columns
184
185
186def MakeRequestUrl(state, url):
187 return 'http%s://%s' % ('s' if state.ssl else '', url)
188
189
Fei Shaobd07c9a2020-06-15 19:04:50 +0800190class ProgressBar:
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800191 SIZE_WIDTH = 11
192 SPEED_WIDTH = 10
193 DURATION_WIDTH = 6
194 PERCENTAGE_WIDTH = 8
195
196 def __init__(self, name):
197 self._start_time = time.time()
198 self._name = name
199 self._size = 0
200 self._width = 0
201 self._name_width = 0
202 self._name_max = 0
203 self._stat_width = 0
204 self._max = 0
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800205 self._CalculateSize()
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800206 self.SetProgress(0)
207
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800208 def _CalculateSize(self):
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800209 self._width = GetTerminalSize()[1] or _DEFAULT_TERMINAL_WIDTH
210 self._name_width = int(self._width * 0.3)
211 self._name_max = self._name_width
212 self._stat_width = self.SIZE_WIDTH + self.SPEED_WIDTH + self.DURATION_WIDTH
213 self._max = (self._width - self._name_width - self._stat_width -
214 self.PERCENTAGE_WIDTH)
215
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800216 def _SizeToHuman(self, size_in_bytes):
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800217 if size_in_bytes < 1024:
218 unit = 'B'
219 value = size_in_bytes
220 elif size_in_bytes < 1024 ** 2:
221 unit = 'KiB'
Yilin Yang14d02a22019-11-01 11:32:03 +0800222 value = size_in_bytes / 1024
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800223 elif size_in_bytes < 1024 ** 3:
224 unit = 'MiB'
Yilin Yang14d02a22019-11-01 11:32:03 +0800225 value = size_in_bytes / (1024 ** 2)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800226 elif size_in_bytes < 1024 ** 4:
227 unit = 'GiB'
Yilin Yang14d02a22019-11-01 11:32:03 +0800228 value = size_in_bytes / (1024 ** 3)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800229 return ' %6.1f %3s' % (value, unit)
230
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800231 def _SpeedToHuman(self, speed_in_bs):
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800232 if speed_in_bs < 1024:
233 unit = 'B'
234 value = speed_in_bs
235 elif speed_in_bs < 1024 ** 2:
236 unit = 'K'
Yilin Yang14d02a22019-11-01 11:32:03 +0800237 value = speed_in_bs / 1024
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800238 elif speed_in_bs < 1024 ** 3:
239 unit = 'M'
Yilin Yang14d02a22019-11-01 11:32:03 +0800240 value = speed_in_bs / (1024 ** 2)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800241 elif speed_in_bs < 1024 ** 4:
242 unit = 'G'
Yilin Yang14d02a22019-11-01 11:32:03 +0800243 value = speed_in_bs / (1024 ** 3)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800244 return ' %6.1f%s/s' % (value, unit)
245
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800246 def _DurationToClock(self, duration):
Yilin Yang14d02a22019-11-01 11:32:03 +0800247 return ' %02d:%02d' % (duration // 60, duration % 60)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800248
249 def SetProgress(self, percentage, size=None):
250 current_width = GetTerminalSize()[1]
251 if self._width != current_width:
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800252 self._CalculateSize()
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800253
254 if size is not None:
255 self._size = size
256
257 elapse_time = time.time() - self._start_time
Yilin Yang14d02a22019-11-01 11:32:03 +0800258 speed = self._size / elapse_time
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800259
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800260 size_str = self._SizeToHuman(self._size)
261 speed_str = self._SpeedToHuman(speed)
262 elapse_str = self._DurationToClock(elapse_time)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800263
264 width = int(self._max * percentage / 100.0)
265 sys.stdout.write(
266 '%*s' % (- self._name_max,
267 self._name if len(self._name) <= self._name_max else
268 self._name[:self._name_max - 4] + ' ...') +
269 size_str + speed_str + elapse_str +
270 ((' [' + '#' * width + ' ' * (self._max - width) + ']' +
271 '%4d%%' % int(percentage)) if self._max > 2 else '') + '\r')
272 sys.stdout.flush()
273
274 def End(self):
275 self.SetProgress(100.0)
276 sys.stdout.write('\n')
277 sys.stdout.flush()
278
279
Fei Shaobd07c9a2020-06-15 19:04:50 +0800280class DaemonState:
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800281 """DaemonState is used for storing Overlord state info."""
282 def __init__(self):
283 self.version_sha1sum = GetVersionDigest()
284 self.host = None
285 self.port = None
286 self.ssl = False
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800287 self.ssl_self_signed = False
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800288 self.ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLSv1_2)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800289 self.ssh = False
290 self.orig_host = None
291 self.ssh_pid = None
292 self.username = None
293 self.password = None
294 self.selected_mid = None
295 self.forwards = {}
296 self.listing = []
297 self.last_list = 0
298
299
Fei Shaobd07c9a2020-06-15 19:04:50 +0800300class OverlordClientDaemon:
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800301 """Overlord Client Daemon."""
302 def __init__(self):
Stimim Chena30447c2020-10-06 10:04:00 +0800303 # Use full module path for jsonrpclib to resolve.
304 import cros.factory.tools.ovl
305 self._state = cros.factory.tools.ovl.DaemonState()
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800306 self._server = None
307
308 def Start(self):
309 self.StartRPCServer()
310
311 def StartRPCServer(self):
312 self._server = SimpleJSONRPCServer(_OVERLORD_CLIENT_DAEMON_RPC_ADDR,
313 logRequests=False)
314 exports = [
315 (self.State, 'State'),
316 (self.Ping, 'Ping'),
317 (self.GetPid, 'GetPid'),
318 (self.Connect, 'Connect'),
319 (self.Clients, 'Clients'),
320 (self.SelectClient, 'SelectClient'),
321 (self.AddForward, 'AddForward'),
322 (self.RemoveForward, 'RemoveForward'),
323 (self.RemoveAllForward, 'RemoveAllForward'),
324 ]
325 for func, name in exports:
326 self._server.register_function(func, name)
327
328 pid = os.fork()
329 if pid == 0:
Peter Shih4d55ded2017-07-03 17:19:01 +0800330 for fd in range(3):
331 os.close(fd)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800332 self._server.serve_forever()
333
334 @staticmethod
335 def GetRPCServer():
336 """Returns the Overlord client daemon RPC server."""
337 server = jsonrpclib.Server('http://%s:%d' %
338 _OVERLORD_CLIENT_DAEMON_RPC_ADDR)
339 try:
340 server.Ping()
341 except Exception:
342 return None
343 return server
344
345 def State(self):
346 return self._state
347
348 def Ping(self):
349 return True
350
351 def GetPid(self):
352 return os.getpid()
353
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800354 def _GetJSON(self, path):
355 url = '%s:%d%s' % (self._state.host, self._state.port, path)
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800356 return json.loads(UrlOpen(self._state, url).read())
357
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800358 def _TLSEnabled(self):
359 """Determine if TLS is enabled on given server address."""
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800360 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
361 try:
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800362 # Allow any certificate since we only want to check if server talks TLS.
363 context = ssl.SSLContext(ssl.PROTOCOL_TLSv1_2)
364 context.verify_mode = ssl.CERT_NONE
365
366 sock = context.wrap_socket(sock, server_hostname=self._state.host)
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800367 sock.settimeout(_CONNECT_TIMEOUT)
368 sock.connect((self._state.host, self._state.port))
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800369 return True
Wei-Ning Huangecc80b82016-07-01 16:55:10 +0800370 except ssl.SSLError:
371 return False
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800372 except socket.error: # Connect refused or timeout
373 raise
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800374 except Exception:
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800375 return False # For whatever reason above failed, assume False
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800376
Wei-Ning Huang4f8f6092017-06-09 19:35:35 +0800377 def _CheckTLSCertificate(self, check_hostname=True):
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800378 """Check TLS certificate.
379
380 Returns:
381 A tupple (check_result, if_certificate_is_loaded)
382 """
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800383 def _DoConnect(context):
384 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
385 try:
386 sock.settimeout(_CONNECT_TIMEOUT)
387 sock = context.wrap_socket(sock, server_hostname=self._state.host)
388 sock.connect((self._state.host, self._state.port))
389 except ssl.SSLError:
390 return False
391 finally:
392 sock.close()
393
394 # Save SSLContext for future use.
395 self._state.ssl_context = context
396 return True
397
398 # First try connect with built-in certificates
399 tls_context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
400 if _DoConnect(tls_context):
401 return True
402
403 # Try with already saved certificate, if any.
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800404 tls_context = ssl.SSLContext(ssl.PROTOCOL_TLSv1_2)
405 tls_context.verify_mode = ssl.CERT_REQUIRED
Wei-Ning Huang4f8f6092017-06-09 19:35:35 +0800406 tls_context.check_hostname = check_hostname
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800407
408 tls_cert_path = GetTLSCertPath(self._state.host)
409 if os.path.exists(tls_cert_path):
410 tls_context.load_verify_locations(tls_cert_path)
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800411 self._state.ssl_self_signed = True
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800412
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800413 return _DoConnect(tls_context)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800414
415 def Connect(self, host, port=_OVERLORD_HTTP_PORT, ssh_pid=None,
Wei-Ning Huang4f8f6092017-06-09 19:35:35 +0800416 username=None, password=None, orig_host=None,
417 check_hostname=True):
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800418 self._state.username = username
419 self._state.password = password
420 self._state.host = host
421 self._state.port = port
422 self._state.ssl = False
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800423 self._state.ssl_self_signed = False
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800424 self._state.orig_host = orig_host
425 self._state.ssh_pid = ssh_pid
426 self._state.selected_mid = None
427
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800428 tls_enabled = self._TLSEnabled()
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800429 if tls_enabled:
Wei-Ning Huang4f8f6092017-06-09 19:35:35 +0800430 result = self._CheckTLSCertificate(check_hostname)
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800431 if not result:
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800432 if self._state.ssl_self_signed:
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800433 return ('SSLCertificateChanged', ssl.get_server_certificate(
434 (self._state.host, self._state.port)))
Yilin Yang15a3f8f2020-01-03 17:49:00 +0800435 return ('SSLVerifyFailed', ssl.get_server_certificate(
436 (self._state.host, self._state.port)))
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800437
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800438 try:
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800439 self._state.ssl = tls_enabled
440 UrlOpen(self._state, '%s:%d' % (host, port))
Yilin Yangf54fb912020-01-08 11:42:38 +0800441 except urllib.error.HTTPError as e:
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800442 return ('HTTPError', e.getcode(), str(e), e.read().strip())
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800443 except Exception as e:
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800444 return str(e)
Wei-Ning Huangba768ab2016-02-07 14:38:06 +0800445 else:
446 return True
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800447
448 def Clients(self):
449 if time.time() - self._state.last_list <= _LIST_CACHE_TIMEOUT:
450 return self._state.listing
451
Peter Shihcf0f3b22017-06-19 15:59:22 +0800452 self._state.listing = self._GetJSON('/api/agents/list')
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800453 self._state.last_list = time.time()
454 return self._state.listing
455
456 def SelectClient(self, mid):
457 self._state.selected_mid = mid
458
459 def AddForward(self, mid, remote, local, pid):
460 self._state.forwards[local] = (mid, remote, pid)
461
462 def RemoveForward(self, local_port):
463 try:
464 unused_mid, unused_remote, pid = self._state.forwards[local_port]
465 KillGraceful(pid)
466 del self._state.forwards[local_port]
467 except (KeyError, OSError):
468 pass
469
470 def RemoveAllForward(self):
471 for unused_mid, unused_remote, pid in self._state.forwards.values():
472 try:
473 KillGraceful(pid)
474 except OSError:
475 pass
476 self._state.forwards = {}
477
478
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800479class SSLEnabledWebSocketBaseClient(WebSocketBaseClient):
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800480 def __init__(self, state, *args, **kwargs):
481 cafile = ssl.get_default_verify_paths().openssl_cafile
482 # For some system / distribution, python can not detect system cafile path.
483 # In such case we fallback to the default path.
484 if not os.path.exists(cafile):
485 cafile = '/etc/ssl/certs/ca-certificates.crt'
486
487 if state.ssl_self_signed:
488 cafile = GetTLSCertPath(state.host)
489
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800490 ssl_options = {
491 'cert_reqs': ssl.CERT_REQUIRED,
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800492 'ca_certs': cafile
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800493 }
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800494 # ws4py does not allow you to specify SSLContext, but rather passing in the
495 # argument of ssl.wrap_socket
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800496 super(SSLEnabledWebSocketBaseClient, self).__init__(
497 ssl_options=ssl_options, *args, **kwargs)
498
499
500class TerminalWebSocketClient(SSLEnabledWebSocketBaseClient):
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800501 def __init__(self, state, mid, escape, *args, **kwargs):
502 super(TerminalWebSocketClient, self).__init__(state, *args, **kwargs)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800503 self._mid = mid
Wei-Ning Huang0c520e92016-03-19 20:01:10 +0800504 self._escape = escape
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800505 self._stdin_fd = sys.stdin.fileno()
506 self._old_termios = None
507
508 def handshake_ok(self):
509 pass
510
511 def opened(self):
512 nonlocals = {'size': (80, 40)}
513
514 def _ResizeWindow():
515 size = GetTerminalSize()
516 if size != nonlocals['size']: # Size not changed, ignore
517 control = {'command': 'resize', 'params': list(size)}
518 payload = chr(_CONTROL_START) + json.dumps(control) + chr(_CONTROL_END)
519 nonlocals['size'] = size
520 try:
521 self.send(payload, binary=True)
522 except Exception:
523 pass
524
525 def _FeedInput():
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800526 self._old_termios = termios.tcgetattr(self._stdin_fd)
527 tty.setraw(self._stdin_fd)
528
529 READY, ENTER_PRESSED, ESCAPE_PRESSED = range(3)
530
531 try:
532 state = READY
533 while True:
Wei-Ning Huang85e763d2016-01-21 15:53:18 +0800534 # Check if terminal is resized
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800535 _ResizeWindow()
536
Wei-Ning Huang85e763d2016-01-21 15:53:18 +0800537 ch = sys.stdin.read(1)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800538
Wei-Ning Huang85e763d2016-01-21 15:53:18 +0800539 # Scan for escape sequence
Wei-Ning Huang0c520e92016-03-19 20:01:10 +0800540 if self._escape:
541 if state == READY:
542 state = ENTER_PRESSED if ch == chr(0x0d) else READY
543 elif state == ENTER_PRESSED:
544 state = ESCAPE_PRESSED if ch == self._escape else READY
545 elif state == ESCAPE_PRESSED:
546 if ch == '.':
547 self.close()
548 break
549 else:
550 state = READY
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800551
Wei-Ning Huang85e763d2016-01-21 15:53:18 +0800552 self.send(ch)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800553 except (KeyboardInterrupt, RuntimeError):
554 pass
555
556 t = threading.Thread(target=_FeedInput)
557 t.daemon = True
558 t.start()
559
560 def closed(self, code, reason=None):
Peter Shihf84a8972017-06-19 15:18:24 +0800561 del code, reason # Unused.
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800562 termios.tcsetattr(self._stdin_fd, termios.TCSANOW, self._old_termios)
563 print('Connection to %s closed.' % self._mid)
564
Yilin Yangf64670b2020-01-06 11:22:18 +0800565 def received_message(self, message):
566 if message.is_binary:
Stimim Chena30447c2020-10-06 10:04:00 +0800567 sys.stdout.write(message.data.decode('utf-8'))
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800568 sys.stdout.flush()
569
570
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800571class ShellWebSocketClient(SSLEnabledWebSocketBaseClient):
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800572 def __init__(self, state, output, *args, **kwargs):
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800573 """Constructor.
574
575 Args:
576 output: output file object.
577 """
578 self.output = output
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800579 super(ShellWebSocketClient, self).__init__(state, *args, **kwargs)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800580
581 def handshake_ok(self):
582 pass
583
584 def opened(self):
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800585 def _FeedInput():
586 try:
587 while True:
Stimim Chena30447c2020-10-06 10:04:00 +0800588 data = sys.stdin.buffer.read(1)
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800589
Peter Shihf84a8972017-06-19 15:18:24 +0800590 if not data:
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800591 self.send(_STDIN_CLOSED * 2)
592 break
593 self.send(data, binary=True)
594 except (KeyboardInterrupt, RuntimeError):
595 pass
596
597 t = threading.Thread(target=_FeedInput)
598 t.daemon = True
599 t.start()
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800600
601 def closed(self, code, reason=None):
602 pass
603
Yilin Yangf64670b2020-01-06 11:22:18 +0800604 def received_message(self, message):
605 if message.is_binary:
Stimim Chena30447c2020-10-06 10:04:00 +0800606 self.output.write(message.data.decode('utf-8'))
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800607 self.output.flush()
608
609
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800610class ForwarderWebSocketClient(SSLEnabledWebSocketBaseClient):
Wei-Ning Huang47c79b82016-05-24 01:24:46 +0800611 def __init__(self, state, sock, *args, **kwargs):
612 super(ForwarderWebSocketClient, self).__init__(state, *args, **kwargs)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800613 self._sock = sock
614 self._stop = threading.Event()
615
616 def handshake_ok(self):
617 pass
618
619 def opened(self):
620 def _FeedInput():
621 try:
622 self._sock.setblocking(False)
623 while True:
624 rd, unused_w, unused_x = select.select([self._sock], [], [], 0.5)
625 if self._stop.is_set():
626 break
627 if self._sock in rd:
628 data = self._sock.recv(_BUFSIZ)
Peter Shihf84a8972017-06-19 15:18:24 +0800629 if not data:
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +0800630 self.close()
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800631 break
632 self.send(data, binary=True)
633 except Exception:
634 pass
635 finally:
636 self._sock.close()
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800637
638 t = threading.Thread(target=_FeedInput)
639 t.daemon = True
640 t.start()
641
642 def closed(self, code, reason=None):
Peter Shihf84a8972017-06-19 15:18:24 +0800643 del code, reason # Unused.
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800644 self._stop.set()
645 sys.exit(0)
646
Yilin Yangf64670b2020-01-06 11:22:18 +0800647 def received_message(self, message):
648 if message.is_binary:
649 self._sock.send(message.data)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800650
651
652def Arg(*args, **kwargs):
653 return (args, kwargs)
654
655
656def Command(command, help_msg=None, args=None):
657 """Decorator for adding argparse parameter for a method."""
658 if args is None:
659 args = []
660 def WrapFunc(func):
Peter Shih13e78c52018-01-23 12:57:07 +0800661 @functools.wraps(func)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800662 def Wrapped(*args, **kwargs):
663 return func(*args, **kwargs)
Peter Shih99b73ec2017-06-16 17:54:15 +0800664 # pylint: disable=protected-access
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800665 Wrapped.__arg_attr = {'command': command, 'help': help_msg, 'args': args}
666 return Wrapped
667 return WrapFunc
668
669
670def ParseMethodSubCommands(cls):
671 """Decorator for a class using the @Command decorator.
672
673 This decorator retrieve command info from each method and append it in to the
674 SUBCOMMANDS class variable, which is later used to construct parser.
675 """
Yilin Yang879fbda2020-05-14 13:52:30 +0800676 for unused_key, method in cls.__dict__.items():
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800677 if hasattr(method, '__arg_attr'):
Peter Shih99b73ec2017-06-16 17:54:15 +0800678 # pylint: disable=protected-access
679 cls.SUBCOMMANDS.append(method.__arg_attr)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800680 return cls
681
682
683@ParseMethodSubCommands
Fei Shaobd07c9a2020-06-15 19:04:50 +0800684class OverlordCLIClient:
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800685 """Overlord command line interface client."""
686
687 SUBCOMMANDS = []
688
689 def __init__(self):
690 self._parser = self._BuildParser()
691 self._selected_mid = None
692 self._server = None
693 self._state = None
Wei-Ning Huang0c520e92016-03-19 20:01:10 +0800694 self._escape = None
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800695
696 def _BuildParser(self):
697 root_parser = argparse.ArgumentParser(prog='ovl')
698 subparsers = root_parser.add_subparsers(help='sub-command')
699
700 root_parser.add_argument('-s', dest='selected_mid', action='store',
701 default=None,
702 help='select target to execute command on')
703 root_parser.add_argument('-S', dest='select_mid_before_action',
704 action='store_true', default=False,
705 help='select target before executing command')
Wei-Ning Huang0c520e92016-03-19 20:01:10 +0800706 root_parser.add_argument('-e', dest='escape', metavar='ESCAPE_CHAR',
707 action='store', default=_ESCAPE, type=str,
708 help='set shell escape character, \'none\' to '
709 'disable escape completely')
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800710
711 for attr in self.SUBCOMMANDS:
712 parser = subparsers.add_parser(attr['command'], help=attr['help'])
713 parser.set_defaults(which=attr['command'])
714 for arg in attr['args']:
715 parser.add_argument(*arg[0], **arg[1])
716
717 return root_parser
718
719 def Main(self):
720 # We want to pass the rest of arguments after shell command directly to the
721 # function without parsing it.
722 try:
723 index = sys.argv.index('shell')
724 except ValueError:
725 args = self._parser.parse_args()
726 else:
727 args = self._parser.parse_args(sys.argv[1:index + 1])
728
729 command = args.which
730 self._selected_mid = args.selected_mid
731
Wei-Ning Huang0c520e92016-03-19 20:01:10 +0800732 if args.escape and args.escape != 'none':
733 self._escape = args.escape[0]
734
Wei-Ning Huang5564eea2016-01-19 14:36:45 +0800735 if command == 'start-server':
736 self.StartServer()
737 return
Fei Shao12ecf382020-06-23 18:32:26 +0800738 if command == 'kill-server':
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800739 self.KillServer()
740 return
741
Wei-Ning Huang5564eea2016-01-19 14:36:45 +0800742 self.CheckDaemon()
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800743 if command == 'status':
744 self.Status()
745 return
Fei Shao12ecf382020-06-23 18:32:26 +0800746 if command == 'connect':
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800747 self.Connect(args)
748 return
749
750 # The following command requires connection to the server
751 self.CheckConnection()
752
753 if args.select_mid_before_action:
754 self.SelectClient(store=False)
755
756 if command == 'select':
757 self.SelectClient(args)
758 elif command == 'ls':
Peter Shih99b73ec2017-06-16 17:54:15 +0800759 self.ListClients(args)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800760 elif command == 'shell':
761 command = sys.argv[sys.argv.index('shell') + 1:]
762 self.Shell(command)
763 elif command == 'push':
764 self.Push(args)
765 elif command == 'pull':
766 self.Pull(args)
767 elif command == 'forward':
768 self.Forward(args)
769
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800770 def _SaveTLSCertificate(self, host, cert_pem):
771 try:
772 os.makedirs(_CERT_DIR)
773 except Exception:
774 pass
775 with open(GetTLSCertPath(host), 'w') as f:
776 f.write(cert_pem)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800777
778 def _HTTPPostFile(self, url, filename, progress=None, user=None, passwd=None):
779 """Perform HTTP POST and upload file to Overlord.
780
781 To minimize the external dependencies, we construct the HTTP post request
782 by ourselves.
783 """
784 url = MakeRequestUrl(self._state, url)
785 size = os.stat(filename).st_size
786 boundary = '-----------%s' % _HTTP_BOUNDARY_MAGIC
787 CRLF = '\r\n'
Yilin Yangf54fb912020-01-08 11:42:38 +0800788 parse = urllib.parse.urlparse(url)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800789
790 part_headers = [
791 '--' + boundary,
792 'Content-Disposition: form-data; name="file"; '
793 'filename="%s"' % os.path.basename(filename),
794 'Content-Type: application/octet-stream',
795 '', ''
796 ]
797 part_header = CRLF.join(part_headers)
798 end_part = CRLF + '--' + boundary + '--' + CRLF
799
800 content_length = len(part_header) + size + len(end_part)
801 if parse.scheme == 'http':
Stimim Chena30447c2020-10-06 10:04:00 +0800802 h = http.client.HTTPConnection(parse.netloc)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800803 else:
Stimim Chena30447c2020-10-06 10:04:00 +0800804 h = http.client.HTTPSConnection(parse.netloc,
805 context=self._state.ssl_context)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800806
807 post_path = url[url.index(parse.netloc) + len(parse.netloc):]
808 h.putrequest('POST', post_path)
809 h.putheader('Content-Length', content_length)
810 h.putheader('Content-Type', 'multipart/form-data; boundary=%s' % boundary)
811
812 if user and passwd:
813 h.putheader(*BasicAuthHeader(user, passwd))
814 h.endheaders()
Stimim Chena30447c2020-10-06 10:04:00 +0800815 h.send(part_header.encode('utf-8'))
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800816
817 count = 0
Stimim Chena30447c2020-10-06 10:04:00 +0800818 with open(filename, 'rb') as f:
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800819 while True:
820 data = f.read(_BUFSIZ)
821 if not data:
822 break
823 count += len(data)
824 if progress:
Yilin Yang14d02a22019-11-01 11:32:03 +0800825 progress(count * 100 // size, count)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800826 h.send(data)
827
Stimim Chena30447c2020-10-06 10:04:00 +0800828 h.send(end_part.encode('utf-8'))
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800829 progress(100)
830
831 if count != size:
832 logging.warning('file changed during upload, upload may be truncated.')
833
Stimim Chena30447c2020-10-06 10:04:00 +0800834 resp = h.getresponse()
835 return resp.status == 200
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800836
Wei-Ning Huang5564eea2016-01-19 14:36:45 +0800837 def CheckDaemon(self):
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800838 self._server = OverlordClientDaemon.GetRPCServer()
839 if self._server is None:
840 print('* daemon not running, starting it now on port %d ... *' %
841 _OVERLORD_CLIENT_DAEMON_PORT)
Wei-Ning Huang5564eea2016-01-19 14:36:45 +0800842 self.StartServer()
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800843
844 self._state = self._server.State()
845 sha1sum = GetVersionDigest()
846
847 if sha1sum != self._state.version_sha1sum:
848 print('ovl server is out of date. killing...')
849 KillGraceful(self._server.GetPid())
Wei-Ning Huang5564eea2016-01-19 14:36:45 +0800850 self.StartServer()
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800851
852 def GetSSHControlFile(self, host):
853 return _SSH_CONTROL_SOCKET_PREFIX + host
854
855 def SSHTunnel(self, user, host, port):
856 """SSH forward the remote overlord server.
857
858 Overlord server may not have port 9000 open to the public network, in such
859 case we can SSH forward the port to localhost.
860 """
861
862 control_file = self.GetSSHControlFile(host)
863 try:
864 os.unlink(control_file)
865 except Exception:
866 pass
867
868 subprocess.Popen([
869 'ssh', '-Nf',
870 '-M', # Enable master mode
871 '-S', control_file,
872 '-L', '9000:localhost:9000',
873 '-p', str(port),
874 '%s%s' % (user + '@' if user else '', host)
875 ]).wait()
876
Yilin Yang83c8f442020-05-05 13:46:51 +0800877 p = process_utils.Spawn([
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800878 'ssh',
879 '-S', control_file,
880 '-O', 'check', host,
Yilin Yang83c8f442020-05-05 13:46:51 +0800881 ], read_stderr=True, ignore_stdout=True)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800882
Yilin Yang83c8f442020-05-05 13:46:51 +0800883 s = re.search(r'pid=(\d+)', p.stderr_data)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800884 if s:
885 return int(s.group(1))
886
887 raise RuntimeError('can not establish ssh connection')
888
889 def CheckConnection(self):
890 if self._state.host is None:
891 raise RuntimeError('not connected to any server, abort')
892
893 try:
894 self._server.Clients()
895 except Exception:
896 raise RuntimeError('remote server disconnected, abort')
897
898 if self._state.ssh_pid is not None:
899 ret = subprocess.Popen(['kill', '-0', str(self._state.ssh_pid)],
900 stdout=subprocess.PIPE,
901 stderr=subprocess.PIPE).wait()
902 if ret != 0:
903 raise RuntimeError('ssh tunnel disconnected, please re-connect')
904
905 def CheckClient(self):
906 if self._selected_mid is None:
907 if self._state.selected_mid is None:
908 raise RuntimeError('No client is selected')
909 self._selected_mid = self._state.selected_mid
910
Peter Shihcf0f3b22017-06-19 15:59:22 +0800911 if not any(client['mid'] == self._selected_mid
912 for client in self._server.Clients()):
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800913 raise RuntimeError('client %s disappeared' % self._selected_mid)
914
915 def CheckOutput(self, command):
916 headers = []
917 if self._state.username is not None and self._state.password is not None:
918 headers.append(BasicAuthHeader(self._state.username,
919 self._state.password))
920
921 scheme = 'ws%s://' % ('s' if self._state.ssl else '')
Yilin Yang8d4f9d02019-11-28 17:12:11 +0800922 sio = StringIO()
Peter Shihe03450b2017-06-14 14:02:08 +0800923 ws = ShellWebSocketClient(
924 self._state, sio, scheme + '%s:%d/api/agent/shell/%s?command=%s' % (
925 self._state.host, self._state.port,
Yilin Yangf54fb912020-01-08 11:42:38 +0800926 urllib.parse.quote(self._selected_mid),
927 urllib.parse.quote(command)),
Peter Shihe03450b2017-06-14 14:02:08 +0800928 headers=headers)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800929 ws.connect()
930 ws.run()
931 return sio.getvalue()
932
933 @Command('status', 'show Overlord connection status')
934 def Status(self):
935 if self._state.host is None:
936 print('Not connected to any host.')
937 else:
938 if self._state.ssh_pid is not None:
939 print('Connected to %s with SSH tunneling.' % self._state.orig_host)
940 else:
941 print('Connected to %s:%d.' % (self._state.host, self._state.port))
942
943 if self._selected_mid is None:
944 self._selected_mid = self._state.selected_mid
945
946 if self._selected_mid is None:
947 print('No client is selected.')
948 else:
949 print('Client %s selected.' % self._selected_mid)
950
951 @Command('connect', 'connect to Overlord server', [
952 Arg('host', metavar='HOST', type=str, default='localhost',
953 help='Overlord hostname/IP'),
954 Arg('port', metavar='PORT', type=int,
955 default=_OVERLORD_HTTP_PORT, help='Overlord port'),
956 Arg('-f', '--forward', dest='ssh_forward', default=False,
957 action='store_true',
958 help='connect with SSH forwarding to the host'),
959 Arg('-p', '--ssh-port', dest='ssh_port', default=22,
960 type=int, help='SSH server port for SSH forwarding'),
961 Arg('-l', '--ssh-login', dest='ssh_login', default='',
962 type=str, help='SSH server login name for SSH forwarding'),
963 Arg('-u', '--user', dest='user', default=None,
964 type=str, help='Overlord HTTP auth username'),
965 Arg('-w', '--passwd', dest='passwd', default=None, type=str,
Wei-Ning Huang4f8f6092017-06-09 19:35:35 +0800966 help='Overlord HTTP auth password'),
967 Arg('-i', '--no-check-hostname', dest='check_hostname',
968 default=True, action='store_false',
969 help='Ignore SSL cert hostname check')])
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +0800970 def Connect(self, args):
971 ssh_pid = None
972 host = args.host
973 orig_host = args.host
974
975 if args.ssh_forward:
976 # Kill previous SSH tunnel
977 self.KillSSHTunnel()
978
979 ssh_pid = self.SSHTunnel(args.ssh_login, args.host, args.ssh_port)
980 host = 'localhost'
981
Wei-Ning Huangba768ab2016-02-07 14:38:06 +0800982 username_provided = args.user is not None
983 password_provided = args.passwd is not None
984 prompt = False
985
Peter Shih533566a2018-09-05 17:48:03 +0800986 for unused_i in range(3): # pylint: disable=too-many-nested-blocks
Wei-Ning Huangba768ab2016-02-07 14:38:06 +0800987 try:
988 if prompt:
989 if not username_provided:
Yilin Yang8cc5dfb2019-10-22 15:58:53 +0800990 args.user = input('Username: ')
Wei-Ning Huangba768ab2016-02-07 14:38:06 +0800991 if not password_provided:
992 args.passwd = getpass.getpass('Password: ')
993
994 ret = self._server.Connect(host, args.port, ssh_pid, args.user,
Wei-Ning Huang4f8f6092017-06-09 19:35:35 +0800995 args.passwd, orig_host,
996 args.check_hostname)
Wei-Ning Huangba768ab2016-02-07 14:38:06 +0800997 if isinstance(ret, list):
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +0800998 if ret[0].startswith('SSL'):
999 cert_pem = ret[1]
1000 fp = GetTLSCertificateSHA1Fingerprint(cert_pem)
1001 fp_text = ':'.join([fp[i:i+2] for i in range(0, len(fp), 2)])
1002
1003 if ret[0] == 'SSLCertificateChanged':
1004 print(_TLS_CERT_CHANGED_WARNING % (fp_text, GetTLSCertPath(host)))
1005 return
Fei Shao12ecf382020-06-23 18:32:26 +08001006 if ret[0] == 'SSLVerifyFailed':
Wei-Ning Huang47c79b82016-05-24 01:24:46 +08001007 print(_TLS_CERT_FAILED_WARNING % (fp_text), end='')
Yilin Yang8cc5dfb2019-10-22 15:58:53 +08001008 response = input()
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +08001009 if response.lower() in ['y', 'ye', 'yes']:
1010 self._SaveTLSCertificate(host, cert_pem)
Wei-Ning Huang47c79b82016-05-24 01:24:46 +08001011 print('TLS host Certificate trusted, you will not be prompted '
1012 'next time.\n')
Wei-Ning Huangba768ab2016-02-07 14:38:06 +08001013 continue
Fei Shaof91ab8f2020-06-23 17:54:03 +08001014 print('connection aborted.')
1015 return
Fei Shao12ecf382020-06-23 18:32:26 +08001016 if ret[0] == 'HTTPError':
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +08001017 code, except_str, body = ret[1:]
1018 if code == 401:
1019 print('connect: %s' % body)
1020 prompt = True
1021 if not username_provided or not password_provided:
1022 continue
Fei Shaof91ab8f2020-06-23 17:54:03 +08001023 break
Fei Shao0e4e2c62020-06-23 18:22:26 +08001024 logging.error('%s; %s', except_str, body)
Wei-Ning Huangba768ab2016-02-07 14:38:06 +08001025
1026 if ret is not True:
1027 print('can not connect to %s: %s' % (host, ret))
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +08001028 else:
1029 print('connection to %s:%d established.' % (host, args.port))
Wei-Ning Huangba768ab2016-02-07 14:38:06 +08001030 except Exception as e:
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +08001031 logging.error(e)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001032 else:
Wei-Ning Huangba768ab2016-02-07 14:38:06 +08001033 break
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001034
Wei-Ning Huang5564eea2016-01-19 14:36:45 +08001035 @Command('start-server', 'start overlord CLI client server')
1036 def StartServer(self):
1037 self._server = OverlordClientDaemon.GetRPCServer()
1038 if self._server is None:
1039 OverlordClientDaemon().Start()
1040 time.sleep(1)
1041 self._server = OverlordClientDaemon.GetRPCServer()
1042 if self._server is not None:
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +08001043 print('* daemon started successfully *\n')
Wei-Ning Huang5564eea2016-01-19 14:36:45 +08001044
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001045 @Command('kill-server', 'kill overlord CLI client server')
1046 def KillServer(self):
1047 self._server = OverlordClientDaemon.GetRPCServer()
1048 if self._server is None:
1049 return
1050
Wei-Ning Huangee7ca8d2015-12-12 05:48:02 +08001051 self._state = self._server.State()
1052
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001053 # Kill SSH Tunnel
1054 self.KillSSHTunnel()
1055
1056 # Kill server daemon
1057 KillGraceful(self._server.GetPid())
1058
1059 def KillSSHTunnel(self):
1060 if self._state.ssh_pid is not None:
1061 KillGraceful(self._state.ssh_pid)
1062
Peter Shihcf0f3b22017-06-19 15:59:22 +08001063 def _FilterClients(self, clients, prop_filters, mid=None):
1064 def _ClientPropertiesMatch(client, key, regex):
1065 try:
1066 return bool(re.search(regex, client['properties'][key]))
1067 except KeyError:
1068 return False
1069
1070 for prop_filter in prop_filters:
1071 key, sep, regex = prop_filter.partition('=')
1072 if not sep:
1073 # The filter doesn't contains =.
1074 raise ValueError('Invalid filter condition %r' % filter)
1075 clients = [c for c in clients if _ClientPropertiesMatch(c, key, regex)]
1076
1077 if mid is not None:
1078 client = next((c for c in clients if c['mid'] == mid), None)
1079 if client:
1080 return [client]
1081 clients = [c for c in clients if c['mid'].startswith(mid)]
1082 return clients
1083
1084 @Command('ls', 'list clients', [
1085 Arg('-f', '--filter', default=[], dest='filters', action='append',
1086 help=('Conditions to filter clients by properties. '
1087 'Should be in form "key=regex", where regex is the regular '
1088 'expression that should be found in the value. '
1089 'Multiple --filter arguments would be ANDed.')),
Peter Shih99b73ec2017-06-16 17:54:15 +08001090 Arg('-v', '--verbose', default=False, action='store_true',
1091 help='Print properties of each client.')
1092 ])
1093 def ListClients(self, args):
Peter Shihcf0f3b22017-06-19 15:59:22 +08001094 clients = self._FilterClients(self._server.Clients(), args.filters)
1095 for client in clients:
1096 if args.verbose:
Peter Shih99b73ec2017-06-16 17:54:15 +08001097 print(yaml.safe_dump(client, default_flow_style=False))
Peter Shihcf0f3b22017-06-19 15:59:22 +08001098 else:
1099 print(client['mid'])
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001100
1101 @Command('select', 'select default client', [
Peter Shihcf0f3b22017-06-19 15:59:22 +08001102 Arg('-f', '--filter', default=[], dest='filters', action='append',
1103 help=('Conditions to filter clients by properties. '
1104 'Should be in form "key=regex", where regex is the regular '
1105 'expression that should be found in the value. '
1106 'Multiple --filter arguments would be ANDed.')),
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001107 Arg('mid', metavar='mid', nargs='?', default=None)])
1108 def SelectClient(self, args=None, store=True):
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001109 mid = args.mid if args is not None else None
Wei-Ning Huang6796a132017-06-22 14:46:32 +08001110 filters = args.filters if args is not None else []
1111 clients = self._FilterClients(self._server.Clients(), filters, mid=mid)
Peter Shihcf0f3b22017-06-19 15:59:22 +08001112
1113 if not clients:
1114 raise RuntimeError('select: client not found')
Fei Shao0e4e2c62020-06-23 18:22:26 +08001115 if len(clients) == 1:
Peter Shihcf0f3b22017-06-19 15:59:22 +08001116 mid = clients[0]['mid']
1117 else:
1118 # This case would not happen when args.mid is specified.
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001119 print('Select from the following clients:')
1120 for i, client in enumerate(clients):
Peter Shihcf0f3b22017-06-19 15:59:22 +08001121 print(' %d. %s' % (i + 1, client['mid']))
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001122
1123 print('\nSelection: ', end='')
1124 try:
Yilin Yang8cc5dfb2019-10-22 15:58:53 +08001125 choice = int(input()) - 1
Peter Shihcf0f3b22017-06-19 15:59:22 +08001126 mid = clients[choice]['mid']
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001127 except ValueError:
1128 raise RuntimeError('select: invalid selection')
1129 except IndexError:
1130 raise RuntimeError('select: selection out of range')
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001131
1132 self._selected_mid = mid
1133 if store:
1134 self._server.SelectClient(mid)
1135 print('Client %s selected' % mid)
1136
1137 @Command('shell', 'open a shell or execute a shell command', [
1138 Arg('command', metavar='CMD', nargs='?', help='command to execute')])
1139 def Shell(self, command=None):
1140 if command is None:
1141 command = []
1142 self.CheckClient()
1143
1144 headers = []
1145 if self._state.username is not None and self._state.password is not None:
1146 headers.append(BasicAuthHeader(self._state.username,
1147 self._state.password))
1148
1149 scheme = 'ws%s://' % ('s' if self._state.ssl else '')
Peter Shihe03450b2017-06-14 14:02:08 +08001150 if command:
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001151 cmd = ' '.join(command)
Peter Shihe03450b2017-06-14 14:02:08 +08001152 ws = ShellWebSocketClient(
1153 self._state, sys.stdout,
1154 scheme + '%s:%d/api/agent/shell/%s?command=%s' % (
1155 self._state.host, self._state.port,
Yilin Yangf54fb912020-01-08 11:42:38 +08001156 urllib.parse.quote(self._selected_mid), urllib.parse.quote(cmd)),
Peter Shihe03450b2017-06-14 14:02:08 +08001157 headers=headers)
1158 else:
1159 ws = TerminalWebSocketClient(
1160 self._state, self._selected_mid, self._escape,
1161 scheme + '%s:%d/api/agent/tty/%s' % (
1162 self._state.host, self._state.port,
Yilin Yangf54fb912020-01-08 11:42:38 +08001163 urllib.parse.quote(self._selected_mid)),
Peter Shihe03450b2017-06-14 14:02:08 +08001164 headers=headers)
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +08001165 try:
1166 ws.connect()
1167 ws.run()
1168 except socket.error as e:
1169 if e.errno == 32: # Broken pipe
1170 pass
1171 else:
1172 raise
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001173
1174 @Command('push', 'push a file or directory to remote', [
Wei-Ning Huang37e61102016-02-07 13:41:07 +08001175 Arg('srcs', nargs='+', metavar='SOURCE'),
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001176 Arg('dst', metavar='DESTINATION')])
1177 def Push(self, args):
1178 self.CheckClient()
1179
Wei-Ning Huangee7ca8d2015-12-12 05:48:02 +08001180 @AutoRetry('push', _RETRY_TIMES)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001181 def _push(src, dst):
1182 src_base = os.path.basename(src)
Wei-Ning Huangee7ca8d2015-12-12 05:48:02 +08001183
1184 # Local file is a link
1185 if os.path.islink(src):
1186 pbar = ProgressBar(src_base)
1187 link_path = os.readlink(src)
1188 self.CheckOutput('mkdir -p %(dirname)s; '
1189 'if [ -d "%(dst)s" ]; then '
1190 'ln -sf "%(link_path)s" "%(dst)s/%(link_name)s"; '
1191 'else ln -sf "%(link_path)s" "%(dst)s"; fi' %
1192 dict(dirname=os.path.dirname(dst),
1193 link_path=link_path, dst=dst,
1194 link_name=src_base))
1195 pbar.End()
1196 return
1197
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001198 mode = '0%o' % (0x1FF & os.stat(src).st_mode)
1199 url = ('%s:%d/api/agent/upload/%s?dest=%s&perm=%s' %
Peter Shihe03450b2017-06-14 14:02:08 +08001200 (self._state.host, self._state.port,
Yilin Yangf54fb912020-01-08 11:42:38 +08001201 urllib.parse.quote(self._selected_mid), dst, mode))
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001202 try:
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +08001203 UrlOpen(self._state, url + '&filename=%s' % src_base)
Yilin Yangf54fb912020-01-08 11:42:38 +08001204 except urllib.error.HTTPError as e:
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001205 msg = json.loads(e.read()).get('error', None)
1206 raise RuntimeError('push: %s' % msg)
1207
1208 pbar = ProgressBar(src_base)
1209 self._HTTPPostFile(url, src, pbar.SetProgress,
1210 self._state.username, self._state.password)
1211 pbar.End()
1212
Wei-Ning Huang37e61102016-02-07 13:41:07 +08001213 def _push_single_target(src, dst):
1214 if os.path.isdir(src):
1215 dst_exists = ast.literal_eval(self.CheckOutput(
1216 'stat %s >/dev/null 2>&1 && echo True || echo False' % dst))
1217 for root, unused_x, files in os.walk(src):
1218 # If destination directory does not exist, we should strip the first
1219 # layer of directory. For example: src_dir contains a single file 'A'
1220 #
1221 # push src_dir dest_dir
1222 #
1223 # If dest_dir exists, the resulting directory structure should be:
1224 # dest_dir/src_dir/A
1225 # If dest_dir does not exist, the resulting directory structure should
1226 # be:
1227 # dest_dir/A
1228 dst_root = root if dst_exists else root[len(src):].lstrip('/')
1229 for name in files:
1230 _push(os.path.join(root, name),
1231 os.path.join(dst, dst_root, name))
1232 else:
1233 _push(src, dst)
1234
1235 if len(args.srcs) > 1:
1236 dst_type = self.CheckOutput('stat \'%s\' --printf \'%%F\' '
1237 '2>/dev/null' % args.dst).strip()
1238 if not dst_type:
1239 raise RuntimeError('push: %s: No such file or directory' % args.dst)
1240 if dst_type != 'directory':
1241 raise RuntimeError('push: %s: Not a directory' % args.dst)
1242
1243 for src in args.srcs:
1244 if not os.path.exists(src):
1245 raise RuntimeError('push: can not stat "%s": no such file or directory'
1246 % src)
1247 if not os.access(src, os.R_OK):
1248 raise RuntimeError('push: can not open "%s" for reading' % src)
1249
1250 _push_single_target(src, args.dst)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001251
1252 @Command('pull', 'pull a file or directory from remote', [
1253 Arg('src', metavar='SOURCE'),
1254 Arg('dst', metavar='DESTINATION', default='.', nargs='?')])
1255 def Pull(self, args):
1256 self.CheckClient()
1257
Wei-Ning Huangee7ca8d2015-12-12 05:48:02 +08001258 @AutoRetry('pull', _RETRY_TIMES)
Peter Shihe6afab32018-09-11 17:16:48 +08001259 def _pull(src, dst, ftype, perm=0o644, link=None):
Wei-Ning Huangee7ca8d2015-12-12 05:48:02 +08001260 try:
1261 os.makedirs(os.path.dirname(dst))
1262 except Exception:
1263 pass
1264
1265 src_base = os.path.basename(src)
1266
1267 # Remote file is a link
1268 if ftype == 'l':
1269 pbar = ProgressBar(src_base)
1270 if os.path.exists(dst):
1271 os.remove(dst)
1272 os.symlink(link, dst)
1273 pbar.End()
1274 return
1275
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001276 url = ('%s:%d/api/agent/download/%s?filename=%s' %
Peter Shihe03450b2017-06-14 14:02:08 +08001277 (self._state.host, self._state.port,
Yilin Yangf54fb912020-01-08 11:42:38 +08001278 urllib.parse.quote(self._selected_mid), urllib.parse.quote(src)))
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001279 try:
Wei-Ning Huangd3a57dc2016-02-09 16:31:04 +08001280 h = UrlOpen(self._state, url)
Yilin Yangf54fb912020-01-08 11:42:38 +08001281 except urllib.error.HTTPError as e:
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001282 msg = json.loads(e.read()).get('error', 'unkown error')
1283 raise RuntimeError('pull: %s' % msg)
1284 except KeyboardInterrupt:
1285 return
1286
Wei-Ning Huangee7ca8d2015-12-12 05:48:02 +08001287 pbar = ProgressBar(src_base)
Stimim Chena30447c2020-10-06 10:04:00 +08001288 with open(dst, 'wb') as f:
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001289 os.fchmod(f.fileno(), perm)
1290 total_size = int(h.headers.get('Content-Length'))
1291 downloaded_size = 0
Wei-Ning Huangee7ca8d2015-12-12 05:48:02 +08001292
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001293 while True:
1294 data = h.read(_BUFSIZ)
Peter Shihf84a8972017-06-19 15:18:24 +08001295 if not data:
Wei-Ning Huangee7ca8d2015-12-12 05:48:02 +08001296 break
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001297 downloaded_size += len(data)
Yilin Yang14d02a22019-11-01 11:32:03 +08001298 pbar.SetProgress(downloaded_size * 100 / total_size,
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001299 downloaded_size)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001300 f.write(data)
1301 pbar.End()
1302
1303 # Use find to get a listing of all files under a root directory. The 'stat'
1304 # command is used to retrieve the filename and it's filemode.
1305 output = self.CheckOutput(
1306 'cd $HOME; '
1307 'stat "%(src)s" >/dev/null && '
Wei-Ning Huangee7ca8d2015-12-12 05:48:02 +08001308 'find "%(src)s" \'(\' -type f -o -type l \')\' '
1309 '-printf \'%%m\t%%p\t%%y\t%%l\n\''
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001310 % {'src': args.src})
1311
1312 # We got error from the stat command
1313 if output.startswith('stat: '):
1314 sys.stderr.write(output)
1315 return
1316
Wei-Ning Huangee7ca8d2015-12-12 05:48:02 +08001317 entries = output.strip('\n').split('\n')
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001318 common_prefix = os.path.dirname(args.src)
1319
1320 if len(entries) == 1:
1321 entry = entries[0]
Wei-Ning Huangee7ca8d2015-12-12 05:48:02 +08001322 perm, src_path, ftype, link = entry.split('\t', -1)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001323 if os.path.isdir(args.dst):
1324 dst = os.path.join(args.dst, os.path.basename(src_path))
1325 else:
1326 dst = args.dst
Wei-Ning Huangee7ca8d2015-12-12 05:48:02 +08001327 _pull(src_path, dst, ftype, int(perm, base=8), link)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001328 else:
1329 if not os.path.exists(args.dst):
1330 common_prefix = args.src
1331
1332 for entry in entries:
Wei-Ning Huangee7ca8d2015-12-12 05:48:02 +08001333 perm, src_path, ftype, link = entry.split('\t', -1)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001334 rel_dst = src_path[len(common_prefix):].lstrip('/')
Wei-Ning Huangee7ca8d2015-12-12 05:48:02 +08001335 _pull(src_path, os.path.join(args.dst, rel_dst), ftype,
1336 int(perm, base=8), link)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001337
1338 @Command('forward', 'forward remote port to local port', [
1339 Arg('--list', dest='list_all', action='store_true', default=False,
1340 help='list all port forwarding sessions'),
1341 Arg('--remove', metavar='LOCAL_PORT', dest='remove', type=int,
1342 default=None,
1343 help='remove port forwarding for local port LOCAL_PORT'),
1344 Arg('--remove-all', dest='remove_all', action='store_true',
1345 default=False, help='remove all port forwarding'),
1346 Arg('remote', metavar='REMOTE_PORT', type=int, nargs='?'),
1347 Arg('local', metavar='LOCAL_PORT', type=int, nargs='?')])
1348 def Forward(self, args):
1349 if args.list_all:
1350 max_len = 10
Peter Shihf84a8972017-06-19 15:18:24 +08001351 if self._state.forwards:
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001352 max_len = max([len(v[0]) for v in self._state.forwards.values()])
1353
1354 print('%-*s %-8s %-8s' % (max_len, 'Client', 'Remote', 'Local'))
1355 for local in sorted(self._state.forwards.keys()):
1356 value = self._state.forwards[local]
1357 print('%-*s %-8s %-8s' % (max_len, value[0], value[1], local))
1358 return
1359
1360 if args.remove_all:
1361 self._server.RemoveAllForward()
1362 return
1363
1364 if args.remove:
1365 self._server.RemoveForward(args.remove)
1366 return
1367
1368 self.CheckClient()
1369
Wei-Ning Huang9083b7c2016-01-26 16:44:11 +08001370 if args.remote is None:
1371 raise RuntimeError('remote port not specified')
1372
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001373 if args.local is None:
1374 args.local = args.remote
1375 remote = int(args.remote)
1376 local = int(args.local)
1377
1378 def HandleConnection(conn):
1379 headers = []
1380 if self._state.username is not None and self._state.password is not None:
1381 headers.append(BasicAuthHeader(self._state.username,
1382 self._state.password))
1383
1384 scheme = 'ws%s://' % ('s' if self._state.ssl else '')
1385 ws = ForwarderWebSocketClient(
Wei-Ning Huang47c79b82016-05-24 01:24:46 +08001386 self._state, conn,
Peter Shihe03450b2017-06-14 14:02:08 +08001387 scheme + '%s:%d/api/agent/forward/%s?port=%d' % (
1388 self._state.host, self._state.port,
Yilin Yangf54fb912020-01-08 11:42:38 +08001389 urllib.parse.quote(self._selected_mid), remote),
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001390 headers=headers)
1391 try:
1392 ws.connect()
1393 ws.run()
1394 except Exception as e:
1395 print('error: %s' % e)
1396 finally:
1397 ws.close()
1398
1399 server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
1400 server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
1401 server.bind(('0.0.0.0', local))
1402 server.listen(5)
1403
1404 pid = os.fork()
1405 if pid == 0:
1406 while True:
1407 conn, unused_addr = server.accept()
1408 t = threading.Thread(target=HandleConnection, args=(conn,))
1409 t.daemon = True
1410 t.start()
1411 else:
1412 self._server.AddForward(self._selected_mid, remote, local, pid)
1413
1414
1415def main():
Wei-Ning Huangba768ab2016-02-07 14:38:06 +08001416 # Setup logging format
1417 logger = logging.getLogger()
Stimim Chena30447c2020-10-06 10:04:00 +08001418 logger.setLevel(logging.DEBUG)
Wei-Ning Huangba768ab2016-02-07 14:38:06 +08001419 handler = logging.StreamHandler()
1420 formatter = logging.Formatter('%(asctime)s %(message)s', '%Y/%m/%d %H:%M:%S')
1421 handler.setFormatter(formatter)
1422 logger.addHandler(handler)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001423
1424 # Add DaemonState to JSONRPC lib classes
Stimim Chena30447c2020-10-06 10:04:00 +08001425 config.DEFAULT.classes.add(DaemonState)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001426
1427 ovl = OverlordCLIClient()
1428 try:
1429 ovl.Main()
1430 except KeyboardInterrupt:
1431 print('Ctrl-C received, abort')
1432 except Exception as e:
Stimim Chena30447c2020-10-06 10:04:00 +08001433 logging.exception('exit with error [%s]', e)
Wei-Ning Huang91aaeed2015-09-24 14:51:56 +08001434
1435
1436if __name__ == '__main__':
1437 main()