| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 1 | #!/usr/bin/python -u |
| 2 | # -*- coding: utf-8 -*- |
| 3 | # |
| 4 | # Copyright 2015 The Chromium OS Authors. All rights reserved. |
| 5 | # Use of this source code is governed by a BSD-style license that can be |
| 6 | # found in the LICENSE file. |
| 7 | |
| 8 | from __future__ import print_function |
| 9 | |
| 10 | import argparse |
| 11 | import ast |
| 12 | import base64 |
| 13 | import fcntl |
| 14 | import hashlib |
| 15 | import httplib |
| 16 | import json |
| 17 | import jsonrpclib |
| 18 | import logging |
| 19 | import os |
| 20 | import re |
| 21 | import select |
| 22 | import signal |
| 23 | import socket |
| 24 | import StringIO |
| 25 | import struct |
| 26 | import subprocess |
| 27 | import sys |
| 28 | import tempfile |
| 29 | import termios |
| 30 | import threading |
| 31 | import time |
| 32 | import tty |
| 33 | import urllib2 |
| 34 | import urlparse |
| 35 | |
| 36 | from jsonrpclib.SimpleJSONRPCServer import SimpleJSONRPCServer |
| 37 | from jsonrpclib.config import Config |
| 38 | from ws4py.client import WebSocketBaseClient |
| 39 | |
| 40 | # Python version >= 2.7.9 enables SSL check by default, bypass it. |
| 41 | try: |
| 42 | import ssl |
| 43 | # pylint: disable=W0212 |
| 44 | ssl._create_default_https_context = ssl._create_unverified_context |
| 45 | except Exception: |
| 46 | pass |
| 47 | |
| 48 | |
| 49 | _ESCAPE = '~' |
| 50 | _BUFSIZ = 8192 |
| 51 | _OVERLORD_PORT = 4455 |
| 52 | _OVERLORD_HTTP_PORT = 9000 |
| 53 | _OVERLORD_CLIENT_DAEMON_PORT = 4488 |
| 54 | _OVERLORD_CLIENT_DAEMON_RPC_ADDR = ('127.0.0.1', _OVERLORD_CLIENT_DAEMON_PORT) |
| 55 | |
| Wei-Ning Huang | ee7ca8d | 2015-12-12 05:48:02 +0800 | [diff] [blame] | 56 | _DEFAULT_HTTP_TIMEOUT = 30 |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 57 | _LIST_CACHE_TIMEOUT = 2 |
| 58 | _DEFAULT_TERMINAL_WIDTH = 80 |
| Wei-Ning Huang | ee7ca8d | 2015-12-12 05:48:02 +0800 | [diff] [blame] | 59 | _RETRY_TIMES = 3 |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 60 | |
| 61 | # echo -n overlord | md5sum |
| 62 | _HTTP_BOUNDARY_MAGIC = '9246f080c855a69012707ab53489b921' |
| 63 | |
| 64 | _CONTROL_START = 128 |
| 65 | _CONTROL_END = 129 |
| 66 | _SSH_CONTROL_SOCKET_PREFIX = os.path.join(tempfile.gettempdir(), |
| 67 | 'ovl-ssh-control-') |
| 68 | |
| 69 | # A string that will always be included in the response of |
| 70 | # GET http://OVERLORD_SERVER:_OVERLORD_HTTP_PORT |
| 71 | _OVERLORD_RESPONSE_KEYWORD = '<html>' |
| 72 | |
| 73 | |
| 74 | def GetVersionDigest(): |
| 75 | """Return the sha1sum of the current executing script.""" |
| 76 | with open(__file__, 'r') as f: |
| 77 | return hashlib.sha1(f.read()).hexdigest() |
| 78 | |
| 79 | |
| 80 | def KillGraceful(pid, wait_secs=1): |
| 81 | """Kill a process gracefully by first sending SIGTERM, wait for some time, |
| 82 | then send SIGKILL to make sure it's killed.""" |
| 83 | try: |
| 84 | os.kill(pid, signal.SIGTERM) |
| 85 | time.sleep(wait_secs) |
| 86 | os.kill(pid, signal.SIGKILL) |
| 87 | except OSError: |
| 88 | pass |
| 89 | |
| 90 | |
| Wei-Ning Huang | ee7ca8d | 2015-12-12 05:48:02 +0800 | [diff] [blame] | 91 | def AutoRetry(action_name, retries): |
| 92 | """Decorator for retry function call.""" |
| 93 | def Wrap(func): |
| 94 | def Loop(*args, **kwargs): |
| Wei-Ning Huang | 5564eea | 2016-01-19 14:36:45 +0800 | [diff] [blame] | 95 | for unused_i in range(retries): |
| Wei-Ning Huang | ee7ca8d | 2015-12-12 05:48:02 +0800 | [diff] [blame] | 96 | try: |
| 97 | func(*args, **kwargs) |
| 98 | except Exception as e: |
| Wei-Ning Huang | 5564eea | 2016-01-19 14:36:45 +0800 | [diff] [blame] | 99 | print('error: %s: %s: retrying ...' % (args[0], e)) |
| Wei-Ning Huang | ee7ca8d | 2015-12-12 05:48:02 +0800 | [diff] [blame] | 100 | else: |
| 101 | break |
| 102 | else: |
| Wei-Ning Huang | 5564eea | 2016-01-19 14:36:45 +0800 | [diff] [blame] | 103 | print('error: failed to %s %s' % (action_name, args[0])) |
| Wei-Ning Huang | ee7ca8d | 2015-12-12 05:48:02 +0800 | [diff] [blame] | 104 | return Loop |
| 105 | return Wrap |
| 106 | |
| 107 | |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 108 | def BasicAuthHeader(user, password): |
| 109 | """Return HTTP basic auth header.""" |
| 110 | credential = base64.b64encode('%s:%s' % (user, password)) |
| 111 | return ('Authorization', 'Basic %s' % credential) |
| 112 | |
| 113 | |
| 114 | def GetTerminalSize(): |
| 115 | """Retrieve terminal window size.""" |
| 116 | ws = struct.pack('HHHH', 0, 0, 0, 0) |
| 117 | ws = fcntl.ioctl(0, termios.TIOCGWINSZ, ws) |
| 118 | lines, columns, unused_x, unused_y = struct.unpack('HHHH', ws) |
| 119 | return lines, columns |
| 120 | |
| 121 | |
| 122 | def MakeRequestUrl(state, url): |
| 123 | return 'http%s://%s' % ('s' if state.ssl else '', url) |
| 124 | |
| 125 | |
| 126 | class ProgressBar(object): |
| 127 | SIZE_WIDTH = 11 |
| 128 | SPEED_WIDTH = 10 |
| 129 | DURATION_WIDTH = 6 |
| 130 | PERCENTAGE_WIDTH = 8 |
| 131 | |
| 132 | def __init__(self, name): |
| 133 | self._start_time = time.time() |
| 134 | self._name = name |
| 135 | self._size = 0 |
| 136 | self._width = 0 |
| 137 | self._name_width = 0 |
| 138 | self._name_max = 0 |
| 139 | self._stat_width = 0 |
| 140 | self._max = 0 |
| 141 | self.CalculateSize() |
| 142 | self.SetProgress(0) |
| 143 | |
| 144 | def CalculateSize(self): |
| 145 | self._width = GetTerminalSize()[1] or _DEFAULT_TERMINAL_WIDTH |
| 146 | self._name_width = int(self._width * 0.3) |
| 147 | self._name_max = self._name_width |
| 148 | self._stat_width = self.SIZE_WIDTH + self.SPEED_WIDTH + self.DURATION_WIDTH |
| 149 | self._max = (self._width - self._name_width - self._stat_width - |
| 150 | self.PERCENTAGE_WIDTH) |
| 151 | |
| 152 | def SizeToHuman(self, size_in_bytes): |
| 153 | if size_in_bytes < 1024: |
| 154 | unit = 'B' |
| 155 | value = size_in_bytes |
| 156 | elif size_in_bytes < 1024 ** 2: |
| 157 | unit = 'KiB' |
| 158 | value = size_in_bytes / 1024.0 |
| 159 | elif size_in_bytes < 1024 ** 3: |
| 160 | unit = 'MiB' |
| 161 | value = size_in_bytes / (1024.0 ** 2) |
| 162 | elif size_in_bytes < 1024 ** 4: |
| 163 | unit = 'GiB' |
| 164 | value = size_in_bytes / (1024.0 ** 3) |
| 165 | return ' %6.1f %3s' % (value, unit) |
| 166 | |
| 167 | def SpeedToHuman(self, speed_in_bs): |
| 168 | if speed_in_bs < 1024: |
| 169 | unit = 'B' |
| 170 | value = speed_in_bs |
| 171 | elif speed_in_bs < 1024 ** 2: |
| 172 | unit = 'K' |
| 173 | value = speed_in_bs / 1024.0 |
| 174 | elif speed_in_bs < 1024 ** 3: |
| 175 | unit = 'M' |
| 176 | value = speed_in_bs / (1024.0 ** 2) |
| 177 | elif speed_in_bs < 1024 ** 4: |
| 178 | unit = 'G' |
| 179 | value = speed_in_bs / (1024.0 ** 3) |
| 180 | return ' %6.1f%s/s' % (value, unit) |
| 181 | |
| 182 | def DurationToClock(self, duration): |
| 183 | return ' %02d:%02d' % (duration / 60, duration % 60) |
| 184 | |
| 185 | def SetProgress(self, percentage, size=None): |
| 186 | current_width = GetTerminalSize()[1] |
| 187 | if self._width != current_width: |
| 188 | self.CalculateSize() |
| 189 | |
| 190 | if size is not None: |
| 191 | self._size = size |
| 192 | |
| 193 | elapse_time = time.time() - self._start_time |
| 194 | speed = self._size / float(elapse_time) |
| 195 | |
| 196 | size_str = self.SizeToHuman(self._size) |
| 197 | speed_str = self.SpeedToHuman(speed) |
| 198 | elapse_str = self.DurationToClock(elapse_time) |
| 199 | |
| 200 | width = int(self._max * percentage / 100.0) |
| 201 | sys.stdout.write( |
| 202 | '%*s' % (- self._name_max, |
| 203 | self._name if len(self._name) <= self._name_max else |
| 204 | self._name[:self._name_max - 4] + ' ...') + |
| 205 | size_str + speed_str + elapse_str + |
| 206 | ((' [' + '#' * width + ' ' * (self._max - width) + ']' + |
| 207 | '%4d%%' % int(percentage)) if self._max > 2 else '') + '\r') |
| 208 | sys.stdout.flush() |
| 209 | |
| 210 | def End(self): |
| 211 | self.SetProgress(100.0) |
| 212 | sys.stdout.write('\n') |
| 213 | sys.stdout.flush() |
| 214 | |
| 215 | |
| 216 | class DaemonState(object): |
| 217 | """DaemonState is used for storing Overlord state info.""" |
| 218 | def __init__(self): |
| 219 | self.version_sha1sum = GetVersionDigest() |
| 220 | self.host = None |
| 221 | self.port = None |
| 222 | self.ssl = False |
| 223 | self.ssh = False |
| 224 | self.orig_host = None |
| 225 | self.ssh_pid = None |
| 226 | self.username = None |
| 227 | self.password = None |
| 228 | self.selected_mid = None |
| 229 | self.forwards = {} |
| 230 | self.listing = [] |
| 231 | self.last_list = 0 |
| 232 | |
| 233 | |
| 234 | class OverlordClientDaemon(object): |
| 235 | """Overlord Client Daemon.""" |
| 236 | def __init__(self): |
| 237 | self._state = DaemonState() |
| 238 | self._server = None |
| 239 | |
| 240 | def Start(self): |
| 241 | self.StartRPCServer() |
| 242 | |
| 243 | def StartRPCServer(self): |
| 244 | self._server = SimpleJSONRPCServer(_OVERLORD_CLIENT_DAEMON_RPC_ADDR, |
| 245 | logRequests=False) |
| 246 | exports = [ |
| 247 | (self.State, 'State'), |
| 248 | (self.Ping, 'Ping'), |
| 249 | (self.GetPid, 'GetPid'), |
| 250 | (self.Connect, 'Connect'), |
| 251 | (self.Clients, 'Clients'), |
| 252 | (self.SelectClient, 'SelectClient'), |
| 253 | (self.AddForward, 'AddForward'), |
| 254 | (self.RemoveForward, 'RemoveForward'), |
| 255 | (self.RemoveAllForward, 'RemoveAllForward'), |
| 256 | ] |
| 257 | for func, name in exports: |
| 258 | self._server.register_function(func, name) |
| 259 | |
| 260 | pid = os.fork() |
| 261 | if pid == 0: |
| 262 | self._server.serve_forever() |
| 263 | |
| 264 | @staticmethod |
| 265 | def GetRPCServer(): |
| 266 | """Returns the Overlord client daemon RPC server.""" |
| 267 | server = jsonrpclib.Server('http://%s:%d' % |
| 268 | _OVERLORD_CLIENT_DAEMON_RPC_ADDR) |
| 269 | try: |
| 270 | server.Ping() |
| 271 | except Exception: |
| 272 | return None |
| 273 | return server |
| 274 | |
| 275 | def State(self): |
| 276 | return self._state |
| 277 | |
| 278 | def Ping(self): |
| 279 | return True |
| 280 | |
| 281 | def GetPid(self): |
| 282 | return os.getpid() |
| 283 | |
| 284 | def _UrlOpen(self, url): |
| 285 | """Wrapper for urllib2.urlopen. |
| 286 | |
| 287 | It selects correct HTTP scheme according to self._stat.ssl and add HTTP |
| 288 | basic auth headers. |
| 289 | """ |
| 290 | url = MakeRequestUrl(self._state, url) |
| 291 | request = urllib2.Request(url) |
| 292 | if self._state.username is not None and self._state.password is not None: |
| 293 | request.add_header(*BasicAuthHeader(self._state.username, |
| 294 | self._state.password)) |
| Wei-Ning Huang | ee7ca8d | 2015-12-12 05:48:02 +0800 | [diff] [blame] | 295 | return urllib2.urlopen(request, timeout=_DEFAULT_HTTP_TIMEOUT) |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 296 | |
| 297 | def _GetJSON(self, path): |
| 298 | url = '%s:%d%s' % (self._state.host, self._state.port, path) |
| 299 | return json.loads(self._UrlOpen(url).read()) |
| 300 | |
| 301 | def Connect(self, host, port=_OVERLORD_HTTP_PORT, ssh_pid=None, |
| 302 | username=None, password=None, orig_host=None): |
| 303 | self._state.username = username |
| 304 | self._state.password = password |
| 305 | self._state.host = host |
| 306 | self._state.port = port |
| 307 | self._state.ssl = False |
| 308 | self._state.orig_host = orig_host |
| 309 | self._state.ssh_pid = ssh_pid |
| 310 | self._state.selected_mid = None |
| 311 | |
| 312 | try: |
| 313 | h = self._UrlOpen('%s:%d' % (host, port)) |
| 314 | # Probably not an HTTP server, try HTTPS |
| 315 | if _OVERLORD_RESPONSE_KEYWORD not in h.read(): |
| 316 | self._state.ssl = True |
| 317 | self._UrlOpen('%s:%d' % (host, port)) |
| 318 | except urllib2.HTTPError as e: |
| 319 | logging.exception(e) |
| 320 | return e.getcode() |
| 321 | except Exception as e: |
| 322 | logging.exception(e) |
| 323 | return str(e) |
| 324 | return True |
| 325 | |
| 326 | def Clients(self): |
| 327 | if time.time() - self._state.last_list <= _LIST_CACHE_TIMEOUT: |
| 328 | return self._state.listing |
| 329 | |
| 330 | mids = [client['mid'] for client in self._GetJSON('/api/agents/list')] |
| 331 | self._state.listing = sorted(list(set(mids))) |
| 332 | self._state.last_list = time.time() |
| 333 | return self._state.listing |
| 334 | |
| 335 | def SelectClient(self, mid): |
| 336 | self._state.selected_mid = mid |
| 337 | |
| 338 | def AddForward(self, mid, remote, local, pid): |
| 339 | self._state.forwards[local] = (mid, remote, pid) |
| 340 | |
| 341 | def RemoveForward(self, local_port): |
| 342 | try: |
| 343 | unused_mid, unused_remote, pid = self._state.forwards[local_port] |
| 344 | KillGraceful(pid) |
| 345 | del self._state.forwards[local_port] |
| 346 | except (KeyError, OSError): |
| 347 | pass |
| 348 | |
| 349 | def RemoveAllForward(self): |
| 350 | for unused_mid, unused_remote, pid in self._state.forwards.values(): |
| 351 | try: |
| 352 | KillGraceful(pid) |
| 353 | except OSError: |
| 354 | pass |
| 355 | self._state.forwards = {} |
| 356 | |
| 357 | |
| 358 | class TerminalWebSocketClient(WebSocketBaseClient): |
| 359 | def __init__(self, mid, *args, **kwargs): |
| 360 | super(TerminalWebSocketClient, self).__init__(*args, **kwargs) |
| 361 | self._mid = mid |
| 362 | self._stdin_fd = sys.stdin.fileno() |
| 363 | self._old_termios = None |
| 364 | |
| 365 | def handshake_ok(self): |
| 366 | pass |
| 367 | |
| 368 | def opened(self): |
| 369 | nonlocals = {'size': (80, 40)} |
| 370 | |
| 371 | def _ResizeWindow(): |
| 372 | size = GetTerminalSize() |
| 373 | if size != nonlocals['size']: # Size not changed, ignore |
| 374 | control = {'command': 'resize', 'params': list(size)} |
| 375 | payload = chr(_CONTROL_START) + json.dumps(control) + chr(_CONTROL_END) |
| 376 | nonlocals['size'] = size |
| 377 | try: |
| 378 | self.send(payload, binary=True) |
| 379 | except Exception: |
| 380 | pass |
| 381 | |
| 382 | def _FeedInput(): |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 383 | self._old_termios = termios.tcgetattr(self._stdin_fd) |
| 384 | tty.setraw(self._stdin_fd) |
| 385 | |
| 386 | READY, ENTER_PRESSED, ESCAPE_PRESSED = range(3) |
| 387 | |
| 388 | try: |
| 389 | state = READY |
| 390 | while True: |
| Wei-Ning Huang | 85e763d | 2016-01-21 15:53:18 +0800 | [diff] [blame^] | 391 | # Check if terminal is resized |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 392 | _ResizeWindow() |
| 393 | |
| Wei-Ning Huang | 85e763d | 2016-01-21 15:53:18 +0800 | [diff] [blame^] | 394 | ch = sys.stdin.read(1) |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 395 | |
| Wei-Ning Huang | 85e763d | 2016-01-21 15:53:18 +0800 | [diff] [blame^] | 396 | # Scan for escape sequence |
| 397 | if state == READY: |
| 398 | state = ENTER_PRESSED if ch == chr(0x0d) else READY |
| 399 | elif state == ENTER_PRESSED: |
| 400 | state = ESCAPE_PRESSED if ch == _ESCAPE else READY |
| 401 | elif state == ESCAPE_PRESSED: |
| 402 | if ch == '.': |
| 403 | self.close() |
| 404 | raise RuntimeError('quit') |
| 405 | else: |
| 406 | state = READY |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 407 | |
| Wei-Ning Huang | 85e763d | 2016-01-21 15:53:18 +0800 | [diff] [blame^] | 408 | self.send(ch) |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 409 | except (KeyboardInterrupt, RuntimeError): |
| 410 | pass |
| 411 | |
| 412 | t = threading.Thread(target=_FeedInput) |
| 413 | t.daemon = True |
| 414 | t.start() |
| 415 | |
| 416 | def closed(self, code, reason=None): |
| 417 | termios.tcsetattr(self._stdin_fd, termios.TCSANOW, self._old_termios) |
| 418 | print('Connection to %s closed.' % self._mid) |
| 419 | |
| 420 | def received_message(self, msg): |
| 421 | if msg.is_binary: |
| 422 | sys.stdout.write(msg.data) |
| 423 | sys.stdout.flush() |
| 424 | |
| 425 | |
| 426 | class ShellWebSocketClient(WebSocketBaseClient): |
| 427 | def __init__(self, output, *args, **kwargs): |
| 428 | """Constructor. |
| 429 | |
| 430 | Args: |
| 431 | output: output file object. |
| 432 | """ |
| 433 | self.output = output |
| 434 | super(ShellWebSocketClient, self).__init__(*args, **kwargs) |
| 435 | |
| 436 | def handshake_ok(self): |
| 437 | pass |
| 438 | |
| 439 | def opened(self): |
| 440 | pass |
| 441 | |
| 442 | def closed(self, code, reason=None): |
| 443 | pass |
| 444 | |
| 445 | def received_message(self, msg): |
| 446 | if msg.is_binary: |
| 447 | self.output.write(msg.data) |
| 448 | self.output.flush() |
| 449 | |
| 450 | |
| 451 | class ForwarderWebSocketClient(WebSocketBaseClient): |
| 452 | def __init__(self, sock, *args, **kwargs): |
| 453 | super(ForwarderWebSocketClient, self).__init__(*args, **kwargs) |
| 454 | self._sock = sock |
| 455 | self._stop = threading.Event() |
| 456 | |
| 457 | def handshake_ok(self): |
| 458 | pass |
| 459 | |
| 460 | def opened(self): |
| 461 | def _FeedInput(): |
| 462 | try: |
| 463 | self._sock.setblocking(False) |
| 464 | while True: |
| 465 | rd, unused_w, unused_x = select.select([self._sock], [], [], 0.5) |
| 466 | if self._stop.is_set(): |
| 467 | break |
| 468 | if self._sock in rd: |
| 469 | data = self._sock.recv(_BUFSIZ) |
| 470 | if len(data) == 0: |
| 471 | break |
| 472 | self.send(data, binary=True) |
| 473 | except Exception: |
| 474 | pass |
| 475 | finally: |
| 476 | self._sock.close() |
| 477 | self.close() |
| 478 | |
| 479 | t = threading.Thread(target=_FeedInput) |
| 480 | t.daemon = True |
| 481 | t.start() |
| 482 | |
| 483 | def closed(self, code, reason=None): |
| 484 | self._stop.set() |
| 485 | sys.exit(0) |
| 486 | |
| 487 | def received_message(self, msg): |
| 488 | if msg.is_binary: |
| 489 | self._sock.send(msg.data) |
| 490 | |
| 491 | |
| 492 | def Arg(*args, **kwargs): |
| 493 | return (args, kwargs) |
| 494 | |
| 495 | |
| 496 | def Command(command, help_msg=None, args=None): |
| 497 | """Decorator for adding argparse parameter for a method.""" |
| 498 | if args is None: |
| 499 | args = [] |
| 500 | def WrapFunc(func): |
| 501 | def Wrapped(*args, **kwargs): |
| 502 | return func(*args, **kwargs) |
| 503 | # pylint: disable=W0212 |
| 504 | Wrapped.__arg_attr = {'command': command, 'help': help_msg, 'args': args} |
| 505 | return Wrapped |
| 506 | return WrapFunc |
| 507 | |
| 508 | |
| 509 | def ParseMethodSubCommands(cls): |
| 510 | """Decorator for a class using the @Command decorator. |
| 511 | |
| 512 | This decorator retrieve command info from each method and append it in to the |
| 513 | SUBCOMMANDS class variable, which is later used to construct parser. |
| 514 | """ |
| 515 | for unused_key, method in cls.__dict__.iteritems(): |
| 516 | if hasattr(method, '__arg_attr'): |
| 517 | cls.SUBCOMMANDS.append(method.__arg_attr) # pylint: disable=W0212 |
| 518 | return cls |
| 519 | |
| 520 | |
| 521 | @ParseMethodSubCommands |
| 522 | class OverlordCLIClient(object): |
| 523 | """Overlord command line interface client.""" |
| 524 | |
| 525 | SUBCOMMANDS = [] |
| 526 | |
| 527 | def __init__(self): |
| 528 | self._parser = self._BuildParser() |
| 529 | self._selected_mid = None |
| 530 | self._server = None |
| 531 | self._state = None |
| 532 | |
| 533 | def _BuildParser(self): |
| 534 | root_parser = argparse.ArgumentParser(prog='ovl') |
| 535 | subparsers = root_parser.add_subparsers(help='sub-command') |
| 536 | |
| 537 | root_parser.add_argument('-s', dest='selected_mid', action='store', |
| 538 | default=None, |
| 539 | help='select target to execute command on') |
| 540 | root_parser.add_argument('-S', dest='select_mid_before_action', |
| 541 | action='store_true', default=False, |
| 542 | help='select target before executing command') |
| 543 | |
| 544 | for attr in self.SUBCOMMANDS: |
| 545 | parser = subparsers.add_parser(attr['command'], help=attr['help']) |
| 546 | parser.set_defaults(which=attr['command']) |
| 547 | for arg in attr['args']: |
| 548 | parser.add_argument(*arg[0], **arg[1]) |
| 549 | |
| 550 | return root_parser |
| 551 | |
| 552 | def Main(self): |
| 553 | # We want to pass the rest of arguments after shell command directly to the |
| 554 | # function without parsing it. |
| 555 | try: |
| 556 | index = sys.argv.index('shell') |
| 557 | except ValueError: |
| 558 | args = self._parser.parse_args() |
| 559 | else: |
| 560 | args = self._parser.parse_args(sys.argv[1:index + 1]) |
| 561 | |
| 562 | command = args.which |
| 563 | self._selected_mid = args.selected_mid |
| 564 | |
| Wei-Ning Huang | 5564eea | 2016-01-19 14:36:45 +0800 | [diff] [blame] | 565 | if command == 'start-server': |
| 566 | self.StartServer() |
| 567 | return |
| 568 | elif command == 'kill-server': |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 569 | self.KillServer() |
| 570 | return |
| 571 | |
| Wei-Ning Huang | 5564eea | 2016-01-19 14:36:45 +0800 | [diff] [blame] | 572 | self.CheckDaemon() |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 573 | if command == 'status': |
| 574 | self.Status() |
| 575 | return |
| 576 | elif command == 'connect': |
| 577 | self.Connect(args) |
| 578 | return |
| 579 | |
| 580 | # The following command requires connection to the server |
| 581 | self.CheckConnection() |
| 582 | |
| 583 | if args.select_mid_before_action: |
| 584 | self.SelectClient(store=False) |
| 585 | |
| 586 | if command == 'select': |
| 587 | self.SelectClient(args) |
| 588 | elif command == 'ls': |
| 589 | self.ListClients() |
| 590 | elif command == 'shell': |
| 591 | command = sys.argv[sys.argv.index('shell') + 1:] |
| 592 | self.Shell(command) |
| 593 | elif command == 'push': |
| 594 | self.Push(args) |
| 595 | elif command == 'pull': |
| 596 | self.Pull(args) |
| 597 | elif command == 'forward': |
| 598 | self.Forward(args) |
| 599 | |
| 600 | def _UrlOpen(self, url): |
| 601 | url = MakeRequestUrl(self._state, url) |
| 602 | request = urllib2.Request(url) |
| 603 | if self._state.username is not None and self._state.password is not None: |
| 604 | request.add_header(*BasicAuthHeader(self._state.username, |
| 605 | self._state.password)) |
| Wei-Ning Huang | ee7ca8d | 2015-12-12 05:48:02 +0800 | [diff] [blame] | 606 | return urllib2.urlopen(request, timeout=_DEFAULT_HTTP_TIMEOUT) |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 607 | |
| 608 | def _HTTPPostFile(self, url, filename, progress=None, user=None, passwd=None): |
| 609 | """Perform HTTP POST and upload file to Overlord. |
| 610 | |
| 611 | To minimize the external dependencies, we construct the HTTP post request |
| 612 | by ourselves. |
| 613 | """ |
| 614 | url = MakeRequestUrl(self._state, url) |
| 615 | size = os.stat(filename).st_size |
| 616 | boundary = '-----------%s' % _HTTP_BOUNDARY_MAGIC |
| 617 | CRLF = '\r\n' |
| 618 | parse = urlparse.urlparse(url) |
| 619 | |
| 620 | part_headers = [ |
| 621 | '--' + boundary, |
| 622 | 'Content-Disposition: form-data; name="file"; ' |
| 623 | 'filename="%s"' % os.path.basename(filename), |
| 624 | 'Content-Type: application/octet-stream', |
| 625 | '', '' |
| 626 | ] |
| 627 | part_header = CRLF.join(part_headers) |
| 628 | end_part = CRLF + '--' + boundary + '--' + CRLF |
| 629 | |
| 630 | content_length = len(part_header) + size + len(end_part) |
| 631 | if parse.scheme == 'http': |
| 632 | h = httplib.HTTP(parse.netloc) |
| 633 | else: |
| 634 | h = httplib.HTTPS(parse.netloc) |
| 635 | |
| 636 | post_path = url[url.index(parse.netloc) + len(parse.netloc):] |
| 637 | h.putrequest('POST', post_path) |
| 638 | h.putheader('Content-Length', content_length) |
| 639 | h.putheader('Content-Type', 'multipart/form-data; boundary=%s' % boundary) |
| 640 | |
| 641 | if user and passwd: |
| 642 | h.putheader(*BasicAuthHeader(user, passwd)) |
| 643 | h.endheaders() |
| 644 | h.send(part_header) |
| 645 | |
| 646 | count = 0 |
| 647 | with open(filename, 'r') as f: |
| 648 | while True: |
| 649 | data = f.read(_BUFSIZ) |
| 650 | if not data: |
| 651 | break |
| 652 | count += len(data) |
| 653 | if progress: |
| 654 | progress(int(count * 100.0 / size), count) |
| 655 | h.send(data) |
| 656 | |
| 657 | h.send(end_part) |
| 658 | progress(100) |
| 659 | |
| 660 | if count != size: |
| 661 | logging.warning('file changed during upload, upload may be truncated.') |
| 662 | |
| 663 | errcode, unused_x, unused_y = h.getreply() |
| 664 | return errcode == 200 |
| 665 | |
| Wei-Ning Huang | 5564eea | 2016-01-19 14:36:45 +0800 | [diff] [blame] | 666 | def CheckDaemon(self): |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 667 | self._server = OverlordClientDaemon.GetRPCServer() |
| 668 | if self._server is None: |
| 669 | print('* daemon not running, starting it now on port %d ... *' % |
| 670 | _OVERLORD_CLIENT_DAEMON_PORT) |
| Wei-Ning Huang | 5564eea | 2016-01-19 14:36:45 +0800 | [diff] [blame] | 671 | self.StartServer() |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 672 | |
| 673 | self._state = self._server.State() |
| 674 | sha1sum = GetVersionDigest() |
| 675 | |
| 676 | if sha1sum != self._state.version_sha1sum: |
| 677 | print('ovl server is out of date. killing...') |
| 678 | KillGraceful(self._server.GetPid()) |
| Wei-Ning Huang | 5564eea | 2016-01-19 14:36:45 +0800 | [diff] [blame] | 679 | self.StartServer() |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 680 | |
| 681 | def GetSSHControlFile(self, host): |
| 682 | return _SSH_CONTROL_SOCKET_PREFIX + host |
| 683 | |
| 684 | def SSHTunnel(self, user, host, port): |
| 685 | """SSH forward the remote overlord server. |
| 686 | |
| 687 | Overlord server may not have port 9000 open to the public network, in such |
| 688 | case we can SSH forward the port to localhost. |
| 689 | """ |
| 690 | |
| 691 | control_file = self.GetSSHControlFile(host) |
| 692 | try: |
| 693 | os.unlink(control_file) |
| 694 | except Exception: |
| 695 | pass |
| 696 | |
| 697 | subprocess.Popen([ |
| 698 | 'ssh', '-Nf', |
| 699 | '-M', # Enable master mode |
| 700 | '-S', control_file, |
| 701 | '-L', '9000:localhost:9000', |
| 702 | '-p', str(port), |
| 703 | '%s%s' % (user + '@' if user else '', host) |
| 704 | ]).wait() |
| 705 | |
| 706 | p = subprocess.Popen([ |
| 707 | 'ssh', |
| 708 | '-S', control_file, |
| 709 | '-O', 'check', host, |
| 710 | ], stderr=subprocess.PIPE) |
| 711 | unused_stdout, stderr = p.communicate() |
| 712 | |
| 713 | s = re.search(r'pid=(\d+)', stderr) |
| 714 | if s: |
| 715 | return int(s.group(1)) |
| 716 | |
| 717 | raise RuntimeError('can not establish ssh connection') |
| 718 | |
| 719 | def CheckConnection(self): |
| 720 | if self._state.host is None: |
| 721 | raise RuntimeError('not connected to any server, abort') |
| 722 | |
| 723 | try: |
| 724 | self._server.Clients() |
| 725 | except Exception: |
| 726 | raise RuntimeError('remote server disconnected, abort') |
| 727 | |
| 728 | if self._state.ssh_pid is not None: |
| 729 | ret = subprocess.Popen(['kill', '-0', str(self._state.ssh_pid)], |
| 730 | stdout=subprocess.PIPE, |
| 731 | stderr=subprocess.PIPE).wait() |
| 732 | if ret != 0: |
| 733 | raise RuntimeError('ssh tunnel disconnected, please re-connect') |
| 734 | |
| 735 | def CheckClient(self): |
| 736 | if self._selected_mid is None: |
| 737 | if self._state.selected_mid is None: |
| 738 | raise RuntimeError('No client is selected') |
| 739 | self._selected_mid = self._state.selected_mid |
| 740 | |
| 741 | if self._selected_mid not in self._server.Clients(): |
| 742 | raise RuntimeError('client %s disappeared' % self._selected_mid) |
| 743 | |
| 744 | def CheckOutput(self, command): |
| 745 | headers = [] |
| 746 | if self._state.username is not None and self._state.password is not None: |
| 747 | headers.append(BasicAuthHeader(self._state.username, |
| 748 | self._state.password)) |
| 749 | |
| 750 | scheme = 'ws%s://' % ('s' if self._state.ssl else '') |
| 751 | sio = StringIO.StringIO() |
| 752 | ws = ShellWebSocketClient(sio, |
| 753 | scheme + '%s:%d/api/agent/shell/%s?command=%s' % |
| 754 | (self._state.host, self._state.port, |
| 755 | self._selected_mid, urllib2.quote(command)), |
| 756 | headers=headers) |
| 757 | ws.connect() |
| 758 | ws.run() |
| 759 | return sio.getvalue() |
| 760 | |
| 761 | @Command('status', 'show Overlord connection status') |
| 762 | def Status(self): |
| 763 | if self._state.host is None: |
| 764 | print('Not connected to any host.') |
| 765 | else: |
| 766 | if self._state.ssh_pid is not None: |
| 767 | print('Connected to %s with SSH tunneling.' % self._state.orig_host) |
| 768 | else: |
| 769 | print('Connected to %s:%d.' % (self._state.host, self._state.port)) |
| 770 | |
| 771 | if self._selected_mid is None: |
| 772 | self._selected_mid = self._state.selected_mid |
| 773 | |
| 774 | if self._selected_mid is None: |
| 775 | print('No client is selected.') |
| 776 | else: |
| 777 | print('Client %s selected.' % self._selected_mid) |
| 778 | |
| 779 | @Command('connect', 'connect to Overlord server', [ |
| 780 | Arg('host', metavar='HOST', type=str, default='localhost', |
| 781 | help='Overlord hostname/IP'), |
| 782 | Arg('port', metavar='PORT', type=int, |
| 783 | default=_OVERLORD_HTTP_PORT, help='Overlord port'), |
| 784 | Arg('-f', '--forward', dest='ssh_forward', default=False, |
| 785 | action='store_true', |
| 786 | help='connect with SSH forwarding to the host'), |
| 787 | Arg('-p', '--ssh-port', dest='ssh_port', default=22, |
| 788 | type=int, help='SSH server port for SSH forwarding'), |
| 789 | Arg('-l', '--ssh-login', dest='ssh_login', default='', |
| 790 | type=str, help='SSH server login name for SSH forwarding'), |
| 791 | Arg('-u', '--user', dest='user', default=None, |
| 792 | type=str, help='Overlord HTTP auth username'), |
| 793 | Arg('-w', '--passwd', dest='passwd', default=None, type=str, |
| 794 | help='Overlord HTTP auth password')]) |
| 795 | def Connect(self, args): |
| 796 | ssh_pid = None |
| 797 | host = args.host |
| 798 | orig_host = args.host |
| 799 | |
| 800 | if args.ssh_forward: |
| 801 | # Kill previous SSH tunnel |
| 802 | self.KillSSHTunnel() |
| 803 | |
| 804 | ssh_pid = self.SSHTunnel(args.ssh_login, args.host, args.ssh_port) |
| 805 | host = 'localhost' |
| 806 | |
| 807 | status = self._server.Connect(host, args.port, ssh_pid, args.user, |
| 808 | args.passwd, orig_host) |
| 809 | if status is not True: |
| 810 | if isinstance(status, int): |
| 811 | if status == 401: |
| 812 | msg = '401 Unauthorized' |
| 813 | else: |
| 814 | msg = 'HTTP %d' % status |
| 815 | else: |
| 816 | msg = status |
| 817 | print('can not connect to %s: %s' % (host, msg)) |
| 818 | |
| Wei-Ning Huang | 5564eea | 2016-01-19 14:36:45 +0800 | [diff] [blame] | 819 | @Command('start-server', 'start overlord CLI client server') |
| 820 | def StartServer(self): |
| 821 | self._server = OverlordClientDaemon.GetRPCServer() |
| 822 | if self._server is None: |
| 823 | OverlordClientDaemon().Start() |
| 824 | time.sleep(1) |
| 825 | self._server = OverlordClientDaemon.GetRPCServer() |
| 826 | if self._server is not None: |
| 827 | print('* daemon started successfully *') |
| 828 | |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 829 | @Command('kill-server', 'kill overlord CLI client server') |
| 830 | def KillServer(self): |
| 831 | self._server = OverlordClientDaemon.GetRPCServer() |
| 832 | if self._server is None: |
| 833 | return |
| 834 | |
| Wei-Ning Huang | ee7ca8d | 2015-12-12 05:48:02 +0800 | [diff] [blame] | 835 | self._state = self._server.State() |
| 836 | |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 837 | # Kill SSH Tunnel |
| 838 | self.KillSSHTunnel() |
| 839 | |
| 840 | # Kill server daemon |
| 841 | KillGraceful(self._server.GetPid()) |
| 842 | |
| 843 | def KillSSHTunnel(self): |
| 844 | if self._state.ssh_pid is not None: |
| 845 | KillGraceful(self._state.ssh_pid) |
| 846 | |
| 847 | @Command('ls', 'list all clients') |
| 848 | def ListClients(self): |
| 849 | for client in self._server.Clients(): |
| 850 | print(client) |
| 851 | |
| 852 | @Command('select', 'select default client', [ |
| 853 | Arg('mid', metavar='mid', nargs='?', default=None)]) |
| 854 | def SelectClient(self, args=None, store=True): |
| 855 | clients = self._server.Clients() |
| 856 | |
| 857 | mid = args.mid if args is not None else None |
| 858 | if mid is None: |
| 859 | print('Select from the following clients:') |
| 860 | for i, client in enumerate(clients): |
| 861 | print(' %d. %s' % (i + 1, client)) |
| 862 | |
| 863 | print('\nSelection: ', end='') |
| 864 | try: |
| 865 | choice = int(raw_input()) - 1 |
| 866 | mid = clients[choice] |
| 867 | except ValueError: |
| 868 | raise RuntimeError('select: invalid selection') |
| 869 | except IndexError: |
| 870 | raise RuntimeError('select: selection out of range') |
| 871 | else: |
| 872 | if mid not in clients: |
| 873 | raise RuntimeError('select: client %s does not exist' % mid) |
| 874 | |
| 875 | self._selected_mid = mid |
| 876 | if store: |
| 877 | self._server.SelectClient(mid) |
| 878 | print('Client %s selected' % mid) |
| 879 | |
| 880 | @Command('shell', 'open a shell or execute a shell command', [ |
| 881 | Arg('command', metavar='CMD', nargs='?', help='command to execute')]) |
| 882 | def Shell(self, command=None): |
| 883 | if command is None: |
| 884 | command = [] |
| 885 | self.CheckClient() |
| 886 | |
| 887 | headers = [] |
| 888 | if self._state.username is not None and self._state.password is not None: |
| 889 | headers.append(BasicAuthHeader(self._state.username, |
| 890 | self._state.password)) |
| 891 | |
| 892 | scheme = 'ws%s://' % ('s' if self._state.ssl else '') |
| 893 | if len(command) == 0: |
| 894 | ws = TerminalWebSocketClient(self._selected_mid, |
| 895 | scheme + '%s:%d/api/agent/tty/%s' % |
| 896 | (self._state.host, self._state.port, |
| 897 | self._selected_mid), headers=headers) |
| 898 | else: |
| 899 | cmd = ' '.join(command) |
| 900 | ws = ShellWebSocketClient(sys.stdout, |
| 901 | scheme + '%s:%d/api/agent/shell/%s?command=%s' % |
| 902 | (self._state.host, self._state.port, |
| 903 | self._selected_mid, urllib2.quote(cmd)), |
| 904 | headers=headers) |
| 905 | ws.connect() |
| 906 | ws.run() |
| 907 | |
| 908 | @Command('push', 'push a file or directory to remote', [ |
| 909 | Arg('src', metavar='SOURCE'), |
| 910 | Arg('dst', metavar='DESTINATION')]) |
| 911 | def Push(self, args): |
| 912 | self.CheckClient() |
| 913 | |
| 914 | if not os.path.exists(args.src): |
| 915 | raise RuntimeError('push: can not stat "%s": no such file or directory' |
| 916 | % args.src) |
| 917 | |
| 918 | if not os.access(args.src, os.R_OK): |
| 919 | raise RuntimeError('push: can not open "%s" for reading' % args.src) |
| 920 | |
| Wei-Ning Huang | ee7ca8d | 2015-12-12 05:48:02 +0800 | [diff] [blame] | 921 | @AutoRetry('push', _RETRY_TIMES) |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 922 | def _push(src, dst): |
| 923 | src_base = os.path.basename(src) |
| Wei-Ning Huang | ee7ca8d | 2015-12-12 05:48:02 +0800 | [diff] [blame] | 924 | |
| 925 | # Local file is a link |
| 926 | if os.path.islink(src): |
| 927 | pbar = ProgressBar(src_base) |
| 928 | link_path = os.readlink(src) |
| 929 | self.CheckOutput('mkdir -p %(dirname)s; ' |
| 930 | 'if [ -d "%(dst)s" ]; then ' |
| 931 | 'ln -sf "%(link_path)s" "%(dst)s/%(link_name)s"; ' |
| 932 | 'else ln -sf "%(link_path)s" "%(dst)s"; fi' % |
| 933 | dict(dirname=os.path.dirname(dst), |
| 934 | link_path=link_path, dst=dst, |
| 935 | link_name=src_base)) |
| 936 | pbar.End() |
| 937 | return |
| 938 | |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 939 | mode = '0%o' % (0x1FF & os.stat(src).st_mode) |
| 940 | url = ('%s:%d/api/agent/upload/%s?dest=%s&perm=%s' % |
| 941 | (self._state.host, self._state.port, self._selected_mid, dst, |
| 942 | mode)) |
| 943 | try: |
| 944 | self._UrlOpen(url + '&filename=%s' % src_base) |
| 945 | except urllib2.HTTPError as e: |
| 946 | msg = json.loads(e.read()).get('error', None) |
| 947 | raise RuntimeError('push: %s' % msg) |
| 948 | |
| 949 | pbar = ProgressBar(src_base) |
| 950 | self._HTTPPostFile(url, src, pbar.SetProgress, |
| 951 | self._state.username, self._state.password) |
| 952 | pbar.End() |
| 953 | |
| 954 | if os.path.isdir(args.src): |
| 955 | dst_exists = ast.literal_eval(self.CheckOutput( |
| 956 | 'stat %s >/dev/null 2>&1 && echo True || echo False' % args.dst)) |
| 957 | for root, unused_x, files in os.walk(args.src): |
| 958 | # If destination directory does not exist, we should strip the first |
| 959 | # layer of directory. For example: src_dir contains a single file 'A' |
| 960 | # |
| 961 | # push src_dir dest_dir |
| 962 | # |
| 963 | # If dest_dir exists, the resulting directory structure should be: |
| 964 | # dest_dir/src_dir/A |
| 965 | # If dest_dir does not exist, the resulting directory structure should |
| 966 | # be: |
| 967 | # dest_dir/A |
| 968 | dst_root = root if dst_exists else root[len(args.src):].lstrip('/') |
| 969 | for name in files: |
| 970 | _push(os.path.join(root, name), |
| 971 | os.path.join(args.dst, dst_root, name)) |
| 972 | else: |
| 973 | _push(args.src, args.dst) |
| 974 | |
| 975 | @Command('pull', 'pull a file or directory from remote', [ |
| 976 | Arg('src', metavar='SOURCE'), |
| 977 | Arg('dst', metavar='DESTINATION', default='.', nargs='?')]) |
| 978 | def Pull(self, args): |
| 979 | self.CheckClient() |
| 980 | |
| Wei-Ning Huang | ee7ca8d | 2015-12-12 05:48:02 +0800 | [diff] [blame] | 981 | @AutoRetry('pull', _RETRY_TIMES) |
| 982 | def _pull(src, dst, ftype, perm=0644, link=None): |
| 983 | try: |
| 984 | os.makedirs(os.path.dirname(dst)) |
| 985 | except Exception: |
| 986 | pass |
| 987 | |
| 988 | src_base = os.path.basename(src) |
| 989 | |
| 990 | # Remote file is a link |
| 991 | if ftype == 'l': |
| 992 | pbar = ProgressBar(src_base) |
| 993 | if os.path.exists(dst): |
| 994 | os.remove(dst) |
| 995 | os.symlink(link, dst) |
| 996 | pbar.End() |
| 997 | return |
| 998 | |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 999 | url = ('%s:%d/api/agent/download/%s?filename=%s' % |
| 1000 | (self._state.host, self._state.port, self._selected_mid, |
| 1001 | urllib2.quote(src))) |
| 1002 | try: |
| 1003 | h = self._UrlOpen(url) |
| 1004 | except urllib2.HTTPError as e: |
| 1005 | msg = json.loads(e.read()).get('error', 'unkown error') |
| 1006 | raise RuntimeError('pull: %s' % msg) |
| 1007 | except KeyboardInterrupt: |
| 1008 | return |
| 1009 | |
| Wei-Ning Huang | ee7ca8d | 2015-12-12 05:48:02 +0800 | [diff] [blame] | 1010 | pbar = ProgressBar(src_base) |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 1011 | with open(dst, 'w') as f: |
| 1012 | os.fchmod(f.fileno(), perm) |
| 1013 | total_size = int(h.headers.get('Content-Length')) |
| 1014 | downloaded_size = 0 |
| Wei-Ning Huang | ee7ca8d | 2015-12-12 05:48:02 +0800 | [diff] [blame] | 1015 | |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 1016 | while True: |
| 1017 | data = h.read(_BUFSIZ) |
| Wei-Ning Huang | ee7ca8d | 2015-12-12 05:48:02 +0800 | [diff] [blame] | 1018 | if len(data) == 0: |
| 1019 | break |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 1020 | downloaded_size += len(data) |
| 1021 | pbar.SetProgress(float(downloaded_size) * 100 / total_size, |
| 1022 | downloaded_size) |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 1023 | f.write(data) |
| 1024 | pbar.End() |
| 1025 | |
| 1026 | # Use find to get a listing of all files under a root directory. The 'stat' |
| 1027 | # command is used to retrieve the filename and it's filemode. |
| 1028 | output = self.CheckOutput( |
| 1029 | 'cd $HOME; ' |
| 1030 | 'stat "%(src)s" >/dev/null && ' |
| Wei-Ning Huang | ee7ca8d | 2015-12-12 05:48:02 +0800 | [diff] [blame] | 1031 | 'find "%(src)s" \'(\' -type f -o -type l \')\' ' |
| 1032 | '-printf \'%%m\t%%p\t%%y\t%%l\n\'' |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 1033 | % {'src': args.src}) |
| 1034 | |
| 1035 | # We got error from the stat command |
| 1036 | if output.startswith('stat: '): |
| 1037 | sys.stderr.write(output) |
| 1038 | return |
| 1039 | |
| Wei-Ning Huang | ee7ca8d | 2015-12-12 05:48:02 +0800 | [diff] [blame] | 1040 | entries = output.strip('\n').split('\n') |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 1041 | common_prefix = os.path.dirname(args.src) |
| 1042 | |
| 1043 | if len(entries) == 1: |
| 1044 | entry = entries[0] |
| Wei-Ning Huang | ee7ca8d | 2015-12-12 05:48:02 +0800 | [diff] [blame] | 1045 | perm, src_path, ftype, link = entry.split('\t', -1) |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 1046 | if os.path.isdir(args.dst): |
| 1047 | dst = os.path.join(args.dst, os.path.basename(src_path)) |
| 1048 | else: |
| 1049 | dst = args.dst |
| Wei-Ning Huang | ee7ca8d | 2015-12-12 05:48:02 +0800 | [diff] [blame] | 1050 | _pull(src_path, dst, ftype, int(perm, base=8), link) |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 1051 | else: |
| 1052 | if not os.path.exists(args.dst): |
| 1053 | common_prefix = args.src |
| 1054 | |
| 1055 | for entry in entries: |
| Wei-Ning Huang | ee7ca8d | 2015-12-12 05:48:02 +0800 | [diff] [blame] | 1056 | perm, src_path, ftype, link = entry.split('\t', -1) |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 1057 | rel_dst = src_path[len(common_prefix):].lstrip('/') |
| Wei-Ning Huang | ee7ca8d | 2015-12-12 05:48:02 +0800 | [diff] [blame] | 1058 | _pull(src_path, os.path.join(args.dst, rel_dst), ftype, |
| 1059 | int(perm, base=8), link) |
| Wei-Ning Huang | 91aaeed | 2015-09-24 14:51:56 +0800 | [diff] [blame] | 1060 | |
| 1061 | @Command('forward', 'forward remote port to local port', [ |
| 1062 | Arg('--list', dest='list_all', action='store_true', default=False, |
| 1063 | help='list all port forwarding sessions'), |
| 1064 | Arg('--remove', metavar='LOCAL_PORT', dest='remove', type=int, |
| 1065 | default=None, |
| 1066 | help='remove port forwarding for local port LOCAL_PORT'), |
| 1067 | Arg('--remove-all', dest='remove_all', action='store_true', |
| 1068 | default=False, help='remove all port forwarding'), |
| 1069 | Arg('remote', metavar='REMOTE_PORT', type=int, nargs='?'), |
| 1070 | Arg('local', metavar='LOCAL_PORT', type=int, nargs='?')]) |
| 1071 | def Forward(self, args): |
| 1072 | if args.list_all: |
| 1073 | max_len = 10 |
| 1074 | if len(self._state.forwards): |
| 1075 | max_len = max([len(v[0]) for v in self._state.forwards.values()]) |
| 1076 | |
| 1077 | print('%-*s %-8s %-8s' % (max_len, 'Client', 'Remote', 'Local')) |
| 1078 | for local in sorted(self._state.forwards.keys()): |
| 1079 | value = self._state.forwards[local] |
| 1080 | print('%-*s %-8s %-8s' % (max_len, value[0], value[1], local)) |
| 1081 | return |
| 1082 | |
| 1083 | if args.remove_all: |
| 1084 | self._server.RemoveAllForward() |
| 1085 | return |
| 1086 | |
| 1087 | if args.remove: |
| 1088 | self._server.RemoveForward(args.remove) |
| 1089 | return |
| 1090 | |
| 1091 | self.CheckClient() |
| 1092 | |
| 1093 | if args.local is None: |
| 1094 | args.local = args.remote |
| 1095 | remote = int(args.remote) |
| 1096 | local = int(args.local) |
| 1097 | |
| 1098 | def HandleConnection(conn): |
| 1099 | headers = [] |
| 1100 | if self._state.username is not None and self._state.password is not None: |
| 1101 | headers.append(BasicAuthHeader(self._state.username, |
| 1102 | self._state.password)) |
| 1103 | |
| 1104 | scheme = 'ws%s://' % ('s' if self._state.ssl else '') |
| 1105 | ws = ForwarderWebSocketClient( |
| 1106 | conn, |
| 1107 | scheme + '%s:%d/api/agent/forward/%s?port=%d' % |
| 1108 | (self._state.host, self._state.port, self._selected_mid, remote), |
| 1109 | headers=headers) |
| 1110 | try: |
| 1111 | ws.connect() |
| 1112 | ws.run() |
| 1113 | except Exception as e: |
| 1114 | print('error: %s' % e) |
| 1115 | finally: |
| 1116 | ws.close() |
| 1117 | |
| 1118 | server = socket.socket(socket.AF_INET, socket.SOCK_STREAM) |
| 1119 | server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) |
| 1120 | server.bind(('0.0.0.0', local)) |
| 1121 | server.listen(5) |
| 1122 | |
| 1123 | pid = os.fork() |
| 1124 | if pid == 0: |
| 1125 | while True: |
| 1126 | conn, unused_addr = server.accept() |
| 1127 | t = threading.Thread(target=HandleConnection, args=(conn,)) |
| 1128 | t.daemon = True |
| 1129 | t.start() |
| 1130 | else: |
| 1131 | self._server.AddForward(self._selected_mid, remote, local, pid) |
| 1132 | |
| 1133 | |
| 1134 | def main(): |
| 1135 | logging.basicConfig(level=logging.INFO) |
| 1136 | |
| 1137 | # Add DaemonState to JSONRPC lib classes |
| 1138 | Config.instance().classes.add(DaemonState) |
| 1139 | |
| 1140 | ovl = OverlordCLIClient() |
| 1141 | try: |
| 1142 | ovl.Main() |
| 1143 | except KeyboardInterrupt: |
| 1144 | print('Ctrl-C received, abort') |
| 1145 | except Exception as e: |
| 1146 | print('error: %s' % e) |
| 1147 | |
| 1148 | |
| 1149 | if __name__ == '__main__': |
| 1150 | main() |