b.liu | e958203 | 2025-04-17 19:18:16 +0800 | [diff] [blame^] | 1 | # Host class |
| 2 | # Copyright (c) 2016, Qualcomm Atheros, Inc. |
| 3 | # |
| 4 | # This software may be distributed under the terms of the BSD license. |
| 5 | # See README for more details. |
| 6 | |
| 7 | import logging |
| 8 | import subprocess |
| 9 | import threading |
| 10 | import tempfile |
| 11 | import os |
| 12 | import traceback |
| 13 | import select |
| 14 | |
| 15 | logger = logging.getLogger() |
| 16 | |
| 17 | def remote_compatible(func): |
| 18 | func.remote_compatible = True |
| 19 | return func |
| 20 | |
| 21 | def execute_thread(command, reply): |
| 22 | cmd = ' '.join(command) |
| 23 | logger.debug("thread run: " + cmd) |
| 24 | err = tempfile.TemporaryFile() |
| 25 | try: |
| 26 | status = 0 |
| 27 | buf = subprocess.check_output(command, stderr=err, bufsize=0).decode() |
| 28 | except subprocess.CalledProcessError as e: |
| 29 | status = e.returncode |
| 30 | err.seek(0) |
| 31 | buf = err.read() |
| 32 | err.close() |
| 33 | |
| 34 | logger.debug("thread cmd: " + cmd) |
| 35 | logger.debug("thread exit status: " + str(status)) |
| 36 | logger.debug("thread exit buf: " + str(buf)) |
| 37 | reply.append(status) |
| 38 | reply.append(buf) |
| 39 | |
| 40 | def gen_reaper_file(conf): |
| 41 | fd, filename = tempfile.mkstemp(dir='/tmp', prefix=conf + '-') |
| 42 | f = os.fdopen(fd, 'w') |
| 43 | |
| 44 | f.write("#!/bin/sh\n") |
| 45 | f.write("name=\"$(basename $0)\"\n") |
| 46 | f.write("echo $$ > /tmp/$name.pid\n") |
| 47 | f.write("exec \"$@\"\n"); |
| 48 | |
| 49 | return filename; |
| 50 | |
| 51 | class Host(): |
| 52 | def __init__(self, host=None, ifname=None, port=None, name="", user="root"): |
| 53 | self.host = host |
| 54 | self.name = name |
| 55 | self.user = user |
| 56 | self.monitors = [] |
| 57 | self.monitor_thread = None |
| 58 | self.logs = [] |
| 59 | self.ifname = ifname |
| 60 | self.port = port |
| 61 | self.dev = None |
| 62 | self.monitor_params = [] |
| 63 | if self.name == "" and host != None: |
| 64 | self.name = host |
| 65 | |
| 66 | def local_execute(self, command): |
| 67 | logger.debug("execute: " + str(command)) |
| 68 | err = tempfile.TemporaryFile() |
| 69 | try: |
| 70 | status = 0 |
| 71 | buf = subprocess.check_output(command, stderr=err) |
| 72 | except subprocess.CalledProcessError as e: |
| 73 | status = e.returncode |
| 74 | err.seek(0) |
| 75 | buf = err.read() |
| 76 | err.close() |
| 77 | |
| 78 | logger.debug("status: " + str(status)) |
| 79 | logger.debug("buf: " + str(buf)) |
| 80 | return status, buf.decode() |
| 81 | |
| 82 | def execute(self, command): |
| 83 | if self.host is None: |
| 84 | return self.local_execute(command) |
| 85 | |
| 86 | cmd = ["ssh", self.user + "@" + self.host, ' '.join(command)] |
| 87 | _cmd = self.name + " execute: " + ' '.join(cmd) |
| 88 | logger.debug(_cmd) |
| 89 | err = tempfile.TemporaryFile() |
| 90 | try: |
| 91 | status = 0 |
| 92 | buf = subprocess.check_output(cmd, stderr=err) |
| 93 | except subprocess.CalledProcessError as e: |
| 94 | status = e.returncode |
| 95 | err.seek(0) |
| 96 | buf = err.read() |
| 97 | err.close() |
| 98 | |
| 99 | logger.debug(self.name + " status: " + str(status)) |
| 100 | logger.debug(self.name + " buf: " + str(buf)) |
| 101 | return status, buf.decode() |
| 102 | |
| 103 | # async execute |
| 104 | def thread_run(self, command, res, use_reaper=True): |
| 105 | if use_reaper: |
| 106 | filename = gen_reaper_file("reaper") |
| 107 | self.send_file(filename, filename) |
| 108 | self.execute(["chmod", "755", filename]) |
| 109 | _command = [filename] + command |
| 110 | else: |
| 111 | filename = "" |
| 112 | _command = command |
| 113 | |
| 114 | if self.host is None: |
| 115 | cmd = _command |
| 116 | else: |
| 117 | cmd = ["ssh", self.user + "@" + self.host, ' '.join(_command)] |
| 118 | _cmd = self.name + " thread_run: " + ' '.join(cmd) |
| 119 | logger.debug(_cmd) |
| 120 | t = threading.Thread(target=execute_thread, name=filename, args=(cmd, res)) |
| 121 | t.start() |
| 122 | return t |
| 123 | |
| 124 | def thread_stop(self, t): |
| 125 | if t.name.find("reaper") == -1: |
| 126 | raise Exception("use_reaper required") |
| 127 | |
| 128 | pid_file = t.name + ".pid" |
| 129 | |
| 130 | if t.is_alive(): |
| 131 | cmd = ["kill `cat " + pid_file + "`"] |
| 132 | self.execute(cmd) |
| 133 | |
| 134 | # try again |
| 135 | self.thread_wait(t, 5) |
| 136 | if t.is_alive(): |
| 137 | cmd = ["kill `cat " + pid_file + "`"] |
| 138 | self.execute(cmd) |
| 139 | |
| 140 | # try with -9 |
| 141 | self.thread_wait(t, 5) |
| 142 | if t.is_alive(): |
| 143 | cmd = ["kill -9 `cat " + pid_file + "`"] |
| 144 | self.execute(cmd) |
| 145 | |
| 146 | self.thread_wait(t, 5) |
| 147 | if t.is_alive(): |
| 148 | raise Exception("thread still alive") |
| 149 | |
| 150 | self.execute(["rm", pid_file]) |
| 151 | self.execute(["rm", t.name]) |
| 152 | self.local_execute(["rm", t.name]) |
| 153 | |
| 154 | def thread_wait(self, t, wait=None): |
| 155 | if wait == None: |
| 156 | wait_str = "infinite" |
| 157 | else: |
| 158 | wait_str = str(wait) + "s" |
| 159 | |
| 160 | logger.debug(self.name + " thread_wait(" + wait_str + "): ") |
| 161 | if t.is_alive(): |
| 162 | t.join(wait) |
| 163 | |
| 164 | def pending(self, s, timeout=0): |
| 165 | [r, w, e] = select.select([s], [], [], timeout) |
| 166 | if r: |
| 167 | return True |
| 168 | return False |
| 169 | |
| 170 | def proc_run(self, command): |
| 171 | filename = gen_reaper_file("reaper") |
| 172 | self.send_file(filename, filename) |
| 173 | self.execute(["chmod", "755", filename]) |
| 174 | _command = [filename] + command |
| 175 | |
| 176 | if self.host: |
| 177 | cmd = ["ssh", self.user + "@" + self.host, ' '.join(_command)] |
| 178 | else: |
| 179 | cmd = _command |
| 180 | |
| 181 | _cmd = self.name + " proc_run: " + ' '.join(cmd) |
| 182 | logger.debug(_cmd) |
| 183 | err = tempfile.TemporaryFile() |
| 184 | proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=err, |
| 185 | bufsize=0) |
| 186 | proc.reaper_file = filename |
| 187 | return proc |
| 188 | |
| 189 | def proc_wait_event(self, proc, events, timeout=10): |
| 190 | if not isinstance(events, list): |
| 191 | raise Exception("proc_wait_event() events not a list") |
| 192 | |
| 193 | logger.debug(self.name + " proc_wait_event: " + ' '.join(events) + " timeout: " + str(timeout)) |
| 194 | start = os.times()[4] |
| 195 | try: |
| 196 | while True: |
| 197 | while self.pending(proc.stdout): |
| 198 | line = proc.stdout.readline() |
| 199 | if not line: |
| 200 | return None |
| 201 | line = line.decode() |
| 202 | logger.debug(line.strip('\n')) |
| 203 | for event in events: |
| 204 | if event in line: |
| 205 | return line |
| 206 | now = os.times()[4] |
| 207 | remaining = start + timeout - now |
| 208 | if remaining <= 0: |
| 209 | break |
| 210 | if not self.pending(proc.stdout, timeout=remaining): |
| 211 | break |
| 212 | except: |
| 213 | logger.debug(traceback.format_exc()) |
| 214 | pass |
| 215 | return None |
| 216 | |
| 217 | def proc_stop(self, proc): |
| 218 | if not proc: |
| 219 | return |
| 220 | |
| 221 | self.execute(["kill `cat " + proc.reaper_file + ".pid`"]) |
| 222 | self.execute(["rm", proc.reaper_file + ".pid"]) |
| 223 | self.execute(["rm", proc.reaper_file]) |
| 224 | self.local_execute(["rm", proc.reaper_file]) |
| 225 | proc.kill() |
| 226 | |
| 227 | def proc_dump(self, proc): |
| 228 | if not proc: |
| 229 | return "" |
| 230 | return proc.stdout.read() |
| 231 | |
| 232 | def execute_and_wait_event(self, command, events, timeout=10): |
| 233 | proc = None |
| 234 | ev = None |
| 235 | |
| 236 | try: |
| 237 | proc = self.proc_run(command) |
| 238 | ev = self.proc_wait_event(proc, events, timeout) |
| 239 | except: |
| 240 | pass |
| 241 | |
| 242 | self.proc_stop(proc) |
| 243 | return ev |
| 244 | |
| 245 | def add_log(self, log_file): |
| 246 | self.logs.append(log_file) |
| 247 | |
| 248 | def get_logs(self, local_log_dir=None): |
| 249 | for log in self.logs: |
| 250 | if local_log_dir: |
| 251 | self.local_execute(["scp", self.user + "@[" + self.host + "]:" + log, local_log_dir]) |
| 252 | self.execute(["rm", log]) |
| 253 | del self.logs[:] |
| 254 | |
| 255 | def send_file(self, src, dst): |
| 256 | if self.host is None: |
| 257 | return |
| 258 | self.local_execute(["scp", src, |
| 259 | self.user + "@[" + self.host + "]:" + dst]) |