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
7import logging
8import subprocess
9import threading
10import tempfile
11import os
12import traceback
13import select
14
15logger = logging.getLogger()
16
17def remote_compatible(func):
18    func.remote_compatible = True
19    return func
20
21def 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
40def 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
51class 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])
260