1#!/usr/bin/env python
2#
3# Copyright (C) 2018 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16#
17
18import datetime
19import logging
20import webapp2
21
22from google.appengine.ext import ndb
23
24from webapp.src import vtslab_status as Status
25from webapp.src.proto import model
26from webapp.src.utils import email_util
27from webapp.src.utils import logger
28
29DEVICE_RESPONSE_TIMEOUT_SECONDS = 300
30
31
32class PeriodicDeviceHeartBeat(webapp2.RequestHandler):
33    """Main class for /tasks/device_heartbeat.
34
35    Used to find lost devices and change their status properly.
36
37    Attributes:
38        logger: Logger class
39    """
40    logger = logger.Logger()
41
42    def get(self):
43        """Generates an HTML page based on the task schedules kept in DB."""
44        self.logger.Clear()
45
46        device_query = model.DeviceModel.query(
47            model.DeviceModel.status !=
48            Status.DEVICE_STATUS_DICT["no-response"])
49        devices = device_query.fetch()
50        lost_devices = [
51            x for x in devices
52            if (datetime.datetime.now() - x.timestamp
53                ).seconds >= DEVICE_RESPONSE_TIMEOUT_SECONDS
54        ]
55        devices_to_put = []
56        labs_to_alert = {}
57        for device in lost_devices:
58            self.logger.Println("Device[{}] is not responding.".format(
59                device.serial))
60            device.status = Status.DEVICE_STATUS_DICT["no-response"]
61            devices_to_put.append(device)
62
63            # sending notification
64            lab_query = model.LabModel.query(
65                model.LabModel.hostname == device.hostname)
66            labs = lab_query.fetch()
67            if labs:
68                lab = labs[0]
69                if lab.name not in labs_to_alert:
70                    labs_to_alert[lab.name] = {}
71                    labs_to_alert[lab.name]["_recipients"] = []
72                if device.hostname not in labs_to_alert[lab.name]:
73                    labs_to_alert[lab.name][device.hostname] = []
74                if lab.owner not in labs_to_alert[lab.name]["_recipients"]:
75                    labs_to_alert[lab.name]["_recipients"].append(lab.owner)
76                labs_to_alert[lab.name]["_recipients"].extend([
77                    x for x in lab.admin
78                    if x not in labs_to_alert[lab.name]["_recipients"]
79                ])
80                labs_to_alert[lab.name][device.hostname].append(device.serial)
81            else:
82                logging.warning(
83                    "Could not find a lab model for hostname {}".format(
84                        device.hostname))
85                continue
86
87        if devices_to_put:
88            ndb.put_multi(devices_to_put)
89        if labs_to_alert:
90            email_util.send_device_notification(labs_to_alert)
91
92        self.response.write(
93            "<pre>\n" + "\n".join(self.logger.Get()) + "\n</pre>")
94