1# Copyright 2017 Google Inc. All rights reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14"""Lab Info APIs implemented using Google Cloud Endpoints."""
15
16import datetime
17import endpoints
18import logging
19
20from google.appengine.ext import ndb
21
22from webapp.src import vtslab_status as Status
23from webapp.src.endpoint import endpoint_base
24from webapp.src.endpoint import host_info
25from webapp.src.proto import model
26
27LAB_INFO_RESOURCE = endpoints.ResourceContainer(model.LabInfoMessage)
28LAB_HOST_INFO_RESOURCE = endpoints.ResourceContainer(model.LabHostInfoMessage)
29
30
31@endpoints.api(name='lab', version='v1')
32class LabInfoApi(endpoint_base.EndpointBase):
33    """Endpoint API for lab_info."""
34
35    @endpoints.method(
36        LAB_INFO_RESOURCE,
37        model.DefaultResponse,
38        path="clear",
39        http_method="POST",
40        name="clear")
41    def clear(self, request):
42        """Clears lab info in DB."""
43        lab_query = model.LabModel.query()
44        existing_labs = lab_query.fetch(keys_only=True)
45        if existing_labs and len(existing_labs) > 0:
46            ndb.delete_multi(existing_labs)
47        return model.DefaultResponse(
48            return_code=model.ReturnCodeMessage.SUCCESS)
49
50    @endpoints.method(
51        LAB_INFO_RESOURCE,
52        model.DefaultResponse,
53        path="set",
54        http_method="POST",
55        name="set")
56    def set(self, request):
57        """Sets the lab info based on `request`."""
58        if "host" in [x.name for x in request.all_fields()]:
59            labs_to_put = []
60            for host in request.host:
61                duplicate_query = model.LabModel.query(
62                    model.LabModel.name == request.name,
63                    model.LabModel.owner == request.owner,
64                    model.LabModel.hostname == host.hostname)
65                duplicates = duplicate_query.fetch()
66                if duplicates:
67                    lab = duplicates[0]
68                else:
69                    lab = model.LabModel()
70                lab.name = request.name
71                lab.owner = request.owner
72                lab.admin = request.admin
73                lab.hostname = host.hostname
74                lab.ip = host.ip
75                lab.script = host.script
76
77                null_device_count = 0
78                devices_to_put = []
79                for config_device in host.device:
80                    if config_device.product == "null":
81                        null_device_count += 1
82                        continue
83                    if config_device.serial and config_device.product:
84                        device_query = model.DeviceModel.query(
85                            model.DeviceModel.serial == config_device.serial)
86                        devices = device_query.fetch()
87                        if devices:
88                            device = devices[0]
89                            if (device.hostname != host.hostname) and (
90                                    device.status !=
91                                    Status.DEVICE_STATUS_DICT["no-response"]):
92                                logging.error(
93                                    "{} is alive in another host.".format(
94                                        config_device.serial))
95                                # TODO: send an alert to lab.admin
96                                continue
97                            if device.hostname == host.hostname and set(
98                                    device.device_equipment) == set(
99                                        config_device.device_equipment):
100                                # no need to update.
101                                continue
102                        else:
103                            device = model.DeviceModel()
104                            device.status = Status.DEVICE_STATUS_DICT[
105                                "no-response"]
106                            device.product = config_device.product
107                            device.serial = config_device.serial
108                            device.hostname = host.hostname
109                            device.scheduling_status = (
110                                Status.DEVICE_SCHEDULING_STATUS_DICT["free"])
111                            device.timestamp = datetime.datetime.now()
112                        device.device_equipment = config_device.device_equipment
113                        devices_to_put.append(device)
114                    else:
115                        logging.error("Lab config does not have device "
116                                      "information correctly; it should "
117                                      "specify device product and serial.")
118                if devices_to_put:
119                    ndb.put_multi(devices_to_put)
120
121                lab.timestamp = datetime.datetime.now()
122                labs_to_put.append(lab)
123
124                if null_device_count > 0:
125                    host_info.AddNullDevices(host.hostname, null_device_count)
126
127            if labs_to_put:
128                ndb.put_multi(labs_to_put)
129
130        return model.DefaultResponse(
131            return_code=model.ReturnCodeMessage.SUCCESS)
132
133    @endpoints.method(
134        LAB_HOST_INFO_RESOURCE,
135        model.DefaultResponse,
136        path="set_version",
137        http_method="POST",
138        name="set_version")
139    def set_version(self, request):
140        """Sets vtslab version of the host <hostname>"""
141        lab_query = model.LabModel.query(
142            model.LabModel.hostname == request.hostname)
143        labs = lab_query.fetch()
144
145        labs_to_put = []
146        for lab in labs:
147            lab.vtslab_version = request.vtslab_version.split(":")[0]
148            labs_to_put.append(lab)
149        if labs_to_put:
150            ndb.put_multi(labs_to_put)
151
152        return model.DefaultResponse(
153            return_code=model.ReturnCodeMessage.SUCCESS)
154
155    @endpoints.method(
156        endpoint_base.GET_REQUEST_RESOURCE,
157        model.LabResponseMessage,
158        path="get",
159        http_method="POST",
160        name="get")
161    def get(self, request):
162        """Gets the labs from datastore."""
163        return_list, more = self.Get(request=request,
164                                     metaclass=model.LabModel,
165                                     message=model.LabMessage)
166
167        return model.LabResponseMessage(labs=return_list, has_next=more)
168
169    @endpoints.method(
170        endpoint_base.COUNT_REQUEST_RESOURCE,
171        model.CountResponseMessage,
172        path="count",
173        http_method="POST",
174        name="count")
175    def count(self, request):
176        """Gets total number of BuildModel entities stored in datastore."""
177        filters = self.CreateFilterList(
178            filter_string=request.filter, metaclass=model.LabModel)
179
180        count = self.Count(metaclass=model.LabModel, filters=filters)
181
182        return model.CountResponseMessage(count=count)
183