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