def _FormatResultAsJson(self, result, format_mode=None): if result is None: return dict(status="OK") if format_mode == JsonMode.PROTO3_JSON_MODE: return json.loads( json_format.MessageToJson(result.AsPrimitiveProto())) elif format_mode == JsonMode.GRR_ROOT_TYPES_STRIPPED_JSON_MODE: result_dict = {} for field, value in result.ListSetFields(): if isinstance( field, (rdf_structs.ProtoDynamicEmbedded, rdf_structs.ProtoEmbedded, rdf_structs.ProtoList)): result_dict[field.name] = api_value_renderers.RenderValue( value) else: result_dict[field.name] = api_value_renderers.RenderValue( value)["value"] return result_dict elif format_mode == JsonMode.GRR_TYPE_STRIPPED_JSON_MODE: rendered_data = api_value_renderers.RenderValue(result) return api_value_renderers.StripTypeInfo(rendered_data) elif format_mode == JsonMode.GRR_JSON_MODE: return api_value_renderers.RenderValue(result) else: raise ValueError("Invalid format_mode: %s" % format_mode)
def HandleCheck(self, method_metadata, args=None, replace=None): """Does regression check for given method, args and a replace function.""" if not replace: raise ValueError("replace can't be None") if self.__class__.api_version == 1: request, prepped_request = self._PrepareV1Request( method_metadata.name, args=args) elif self.__class__.api_version == 2: request, prepped_request = self._PrepareV2Request( method_metadata.name, args=args) else: raise ValueError("api_version may be only 1 or 2, not %d" % flags.FLAGS.api_version) session = requests.Session() response = session.send(prepped_request) check_result = { "url": replace(prepped_request.path_url), "method": request.method } if request.data: if compatibility.PY2: data = request.data.decode("utf-8") else: data = request.data request_payload = self._ParseJSON(replace(data)) if request_payload: check_result["request_payload"] = request_payload if (method_metadata.result_type == api_call_router. RouterMethodMetadata.BINARY_STREAM_RESULT_TYPE): check_result["response"] = replace( utils.SmartUnicode(response.content)) else: content = response.content.decode("utf-8") check_result["response"] = self._ParseJSON(replace(content)) if self.__class__.api_version == 1: stripped_response = api_value_renderers.StripTypeInfo( check_result["response"]) if stripped_response != check_result["response"]: check_result["type_stripped_response"] = stripped_response return check_result
def _PrepareV1Request(self, method, args=None): """Prepares API v1 request for a given method and args.""" args_proto = None if args: args_proto = args.AsPrimitiveProto() request = self.connector.BuildRequest(method, args_proto) request.url = request.url.replace("/api/v2/", "/api/") if args and request.data: body_proto = args.__class__().AsPrimitiveProto() json_format.Parse(request.data, body_proto) body_args = args.__class__() body_args.ParseFromString(body_proto.SerializeToString()) request.data = json.dumps( api_value_renderers.StripTypeInfo( api_value_renderers.RenderValue(body_args)), cls=http_api.JSONEncoderWithRDFPrimitivesSupport) prepped_request = request.prepare() return request, prepped_request