1# -*- coding: utf-8 -*-
2#!/usr/bin/env python2.7
3# Copyright (c) 2019 The Chromium OS Authors. All rights reserved.
4# Use of this source code is governed by a BSD-style license that can be
5# found in the LICENSE file.
6
7import json
8import logging
9import requests
10
11"""This class consists of all the helper methods needed to interact with the
12   Datastore @ https://chaos-188802.appspot.com/ used for ChromeOS Interop
13   testing.
14"""
15
16class ChaosDataStoreUtils(object):
17
18    CHAOS_DATASTORE_URL = 'https://chaos-188802.appspot.com'
19
20    # The Datastore defines the following paths for operating methods.
21    ADD_DEVICE = "devices/new"
22    REMOVE_DEVICE = "devices/delete"
23    LOCK_DEVICE = "devices/lock"
24    UNLOCK_DEVICE = 'devices/unlock'
25    SHOW_DEVICE = "devices/"
26    GET_DEVICES = 'devices/'
27    GET_UNLOCKED_DEVICES = "unlocked_devices/"
28    GET_DEVICES_BY_AP_LABEL = "devices/location"
29
30    # HTTP content type. JSON encoded with UTF-8 character encoding.
31    HTTP_HEADER = {'content-type': 'application/json'}
32
33
34    def add_device(self, host_name, ap_label):
35        """
36        Add a device(AP or Packet Capturer) in datastore.
37
38        @param host_name: string, hostname of the device.
39        @param ap_label: string, CrOS_AP (for AP), CrOS_PCAP (for PCAP)
40        @param lab_label: string, CrOS_Chaos (lab name), used for all APs & PCAPs
41
42        @return: True if device was added successfully; False otherwise.
43        @rtype: bool
44
45        """
46        request = self.CHAOS_DATASTORE_URL + '/' + self.ADD_DEVICE
47        logging.debug("Request = %s", request)
48        response = requests.post(request,
49                                 headers=self.HTTP_HEADER,
50                                 data=json.dumps({"hostname":host_name,
51                                                  "ap_label":ap_label,
52                                                  "lab_label":"CrOS_Chaos",
53                                                  "router_name":host_name}))
54        if response.json()['result']:
55            logging.info("Added device %s to datastore", host_name)
56            return True
57
58        return False
59
60
61    def remove_device(self, host_name):
62        """
63        Delete a device(AP or Packet Capturer) in datastore.
64
65        @param host_name: string, hostname of the device to delete.
66
67        @return: True if device was deleted successfully; False otherwise.
68        @rtype: bool
69
70        """
71        request = self.CHAOS_DATASTORE_URL + '/' + self.REMOVE_DEVICE
72        logging.debug("Request = %s", request)
73        response = requests.put(request,
74                                headers=self.HTTP_HEADER,
75                                data=json.dumps({"hostname":host_name}))
76        result_str = "%s deleted." % host_name
77        if result_str in response.text:
78            logging.info("Removed device %s from datastore", host_name)
79            return True
80
81        return False
82
83
84    def lock_device(self, host_name, lock_reason):
85        """
86        Lock a device(AP or Packet Capturer) in datastore.
87
88        @param host_name: string, hostname of the device in datastore.
89
90        @return: True if operation was successful; False otherwise.
91        @rtype: bool
92
93        """
94        request = self.CHAOS_DATASTORE_URL + '/' + self.LOCK_DEVICE
95        logging.debug("Request = %s", request)
96        response = requests.put(request,
97                                headers=self.HTTP_HEADER,
98                                data=json.dumps({"hostname":host_name,
99                                                 "locked_by":lock_reason}))
100        if response.json()['result']:
101            logging.info("Locked device %s in datastore", host_name)
102            return True
103
104        return False
105
106
107    def unlock_device(self, host_name):
108        """
109        Un-lock a device(AP or Packet Capturer) in datastore.
110
111        @param host_name: string, hostname of the device in datastore.
112
113        @return: True if operation was successful; False otherwise.
114        @rtype: bool
115
116        """
117        request = self.CHAOS_DATASTORE_URL + '/' + self.UNLOCK_DEVICE
118        logging.debug("Request = %s", request)
119        response = requests.put(request,
120                                headers=self.HTTP_HEADER,
121                                data=json.dumps({"hostname":host_name}))
122        if response.json()['result']:
123            logging.info("Finished un-locking AP %s in datastore", host_name)
124            return True
125
126        logging.error("Unable to unlock AP %s", host_name)
127        return False
128
129
130    def show_device(self, host_name):
131        """
132        Show device properties for a given device(AP or Packet Capturer).
133
134        @param host_name: string, hostname of the device in datastore to fetch info.
135
136        @return: dict of device name:value properties if successful;
137                 False otherwise.
138        @rtype: dict when True, else bool:False
139
140        """
141        request = self.CHAOS_DATASTORE_URL + '/' + self.SHOW_DEVICE + host_name
142        logging.debug("Request = %s", request)
143        response = requests.get(request)
144        if 'error' in response.text:
145            return False
146
147        return response.json()
148
149
150    def get_unlocked_devices(self):
151        """
152        Get a list of all un-locked devices in the datastore.
153
154        @return: dict of all un-locked devices' name:value properties if successful;
155                 False otherwise.
156        @rtype: dict when True, else bool:False
157
158        """
159        request = self.CHAOS_DATASTORE_URL + '/' + self.GET_UNLOCKED_DEVICES
160        logging.debug("Request = %s", request)
161        response = requests.get(request)
162        if 'error' in response.text:
163            return False
164
165        return response.json()
166
167
168    def get_devices(self):
169        """
170        Get a list of all devices in the datastore.
171
172        @return: dict of all devices' name:value properties if successful;
173                 False otherwise.
174        @rtype: dict when True, else bool:False
175
176        """
177        request = self.CHAOS_DATASTORE_URL + '/' + self.GET_DEVICES
178        logging.debug("Request = %s", request)
179        response = requests.get(request)
180        if 'error' in response.text:
181            return False
182
183        return response.json()
184
185
186    def get_devices_by_type(self, ap_label, lab_label):
187        """
188        Get list of all un-locked devices by ap_label & lab_label
189
190        @param ap_label: string, CrOS_AP/CrOS_PCAP, to filter device types.
191        @param lab_label: string, "CrOS_Chaos", All devices in ChromeOS Chaos lab
192
193        @return: dict of all devices' name:value properties if successful;
194                 False otherwise.
195        @rtype: dict when True, else bool:False
196
197        """
198        request = self.CHAOS_DATASTORE_URL + '/' +  self.GET_DEVICES_BY_AP_LABEL
199        logging.debug("Request = %s", request)
200        response = requests.put(request,
201                                headers=self.HTTP_HEADER,
202                                data=json.dumps({"ap_label":ap_label,
203                                                 "lab_label":lab_label}))
204        if 'error' in response.text:
205            return False
206
207        return response.json()
208
209
210    def find_device(self, host_name):
211        """
212        Find if given device(AP or Packet Capturer) in DataStore.
213
214        @param host_name: string, hostname of the device in datastore to fetch info.
215        @return: True if found; False otherwise.
216        @rtype: bool
217
218        """
219        request = self.CHAOS_DATASTORE_URL + '/' + self.SHOW_DEVICE + host_name
220        logging.debug("Request = %s", request)
221        response = requests.get(request)
222        if 'null' in response.text:
223            return False
224
225        return True
226