def get_signed_cert( common_name: str, san_dns_names: Optional[List[str]] = None, san_ip_addresses: Optional[List[str]] = None, resotocore_uri: str = None, psk: str = None, ca_cert_path: str = None, ) -> Tuple[RSAPrivateKey, Certificate]: if resotocore_uri is None: resotocore_uri = getattr(ArgumentParser.args, "resotocore_uri", None) if psk is None: psk = getattr(ArgumentParser.args, "psk", None) cert_key = gen_rsa_key() cert_csr = gen_csr(cert_key, common_name, san_dns_names, san_ip_addresses) cert_csr_bytes = csr_to_bytes(cert_csr) headers = {} if psk is not None: encode_jwt_to_headers(headers, {}, psk) request_kwargs = {} if ca_cert_path is not None: request_kwargs["verify"] = ca_cert_path r = requests.post(f"{resotocore_uri}/ca/sign", cert_csr_bytes, headers=headers, **request_kwargs) cert_bytes = r.content cert_crt = load_cert_from_bytes(cert_bytes) return cert_key, cert_crt
def connect(self) -> None: resotocore_ws_uri_split = urlsplit(self.resotocore_ws_uri) scheme = resotocore_ws_uri_split.scheme netloc = resotocore_ws_uri_split.netloc path = resotocore_ws_uri_split.path + "/work/queue" query_dict = {"task": ",".join(self.tasks)} query_dict.update( {k: ",".join(v) for k, v in self.task_queue_filter.items()}) query = urlencode(query_dict) ws_uri = urlunsplit((scheme, netloc, path, query, "")) log.debug(f"{self.identifier} connecting to {ws_uri}") headers = {} if getattr(ArgumentParser.args, "psk", None): encode_jwt_to_headers(headers, {}, ArgumentParser.args.psk) self.ws = websocket.WebSocketApp( ws_uri, header=headers, on_open=self.on_open, on_message=self.on_message, on_error=self.on_error, on_close=self.on_close, on_ping=self.on_ping, on_pong=self.on_pong, ) sslopt = None if self.tls_data: sslopt = {"ca_certs": self.tls_data.ca_cert_path} self.ws.run_forever(sslopt=sslopt, ping_interval=30, ping_timeout=10, ping_payload="ping")
def send_graph( graph_export_iterator: GraphExportIterator, resotocore_base_uri: str, resotocore_graph: str, ): merge_uri = f"{resotocore_base_uri}/graph/{resotocore_graph}/merge" log.debug(f"Sending graph via {merge_uri}") headers = { "Content-Type": "application/x-ndjson", "Resoto-Worker-Nodes": str(graph_export_iterator.number_of_nodes), "Resoto-Worker-Edges": str(graph_export_iterator.number_of_edges), } if getattr(ArgumentParser.args, "psk", None): encode_jwt_to_headers(headers, {}, ArgumentParser.args.psk) r = requests.post( merge_uri, data=graph_export_iterator, headers=headers, ) if r.status_code != 200: log.error(r.content) raise RuntimeError(f"Failed to send graph: {r.content}") log.debug(f"resotocore reply: {r.content.decode()}") log.debug(f"Sent {graph_export_iterator.total_lines} items to resotocore")
def post_request( data: Union[bytes, Dict[str, str]], content_type: str, ) -> Optional[Response]: # define content-type headers["Content-Type"] = content_type # sign request if ArgumentParser.args.psk: encode_jwt_to_headers(headers, {}, ArgumentParser.args.psk) body: Optional[Any] = None if isinstance(data, bytes): body = data elif isinstance(data, dict): parts = { name: (name, open(path, "rb"), "application/octet-stream") for name, path in data.items() } body = MultipartEncoder(parts, "file-upload") else: raise AttributeError( f"Can not handle data of type: {type(data)}") try: return post(self.execute_endpoint, data=body, headers=headers, stream=True) except ConnectionError: err = ("Error: Could not communicate with resotocore" f" at {urlsplit(self.execute_endpoint).netloc}." " Is it up and reachable?") print(err, file=sys.stderr)
def update_model(graph: Graph, resotocore_base_uri: str, dump_json: bool = False, tempdir: str = None) -> None: model_uri = f"{resotocore_base_uri}/model" log.debug(f"Updating model via {model_uri}") model_json = json.dumps(graph.export_model(), indent=4) if dump_json: ts = datetime.now().strftime("%Y-%m-%d-%H-%M") with tempfile.NamedTemporaryFile( prefix=f"resoto-model-{ts}-", suffix=".json", delete=not dump_json, dir=tempdir, ) as model_outfile: log.info(f"Writing model json to file {model_outfile.name}") model_outfile.write(model_json.encode()) headers = { "Content-Type": "application/json", } if getattr(ArgumentParser.args, "psk", None): encode_jwt_to_headers(headers, {}, ArgumentParser.args.psk) r = requests.patch(model_uri, data=model_json, headers=headers) if r.status_code != 200: log.error(r.content) raise RuntimeError(f"Failed to create model: {r.content}")
def default_args(resotocore_uri: str = None, psk: str = None, headers=None) -> Tuple[str, str, Dict[str, str]]: if resotocore_uri is None: resotocore_uri = resotocore.http_uri if psk is None: psk = getattr(ArgumentParser.args, "psk", None) if headers is None: headers = {} if psk is not None: encode_jwt_to_headers(headers, {}, psk) return resotocore_uri, psk, headers
def patch_nodes(self, graph: Graph): headers = {"Content-Type": "application/x-ndjson"} if getattr(ArgumentParser.args, "psk", None): encode_jwt_to_headers(headers, {}, ArgumentParser.args.psk) r = requests.patch( f"{self.graph_uri}/nodes", data=GraphChangeIterator(graph), headers=headers ) if r.status_code != 200: err = r.content.decode("utf-8") log.error(err) raise RuntimeError(f"Failed to patch nodes: {err}")
def create_graph(resotocore_base_uri: str, resotocore_graph: str): graph_uri = f"{resotocore_base_uri}/graph/{resotocore_graph}" log.debug(f"Creating graph {resotocore_graph} via {graph_uri}") headers = { "accept": "application/json", "Content-Type": "application/json", } if getattr(ArgumentParser.args, "psk", None): encode_jwt_to_headers(headers, {}, ArgumentParser.args.psk) r = requests.post(graph_uri, data="", headers=headers) if r.status_code != 200: log.error(r.content) raise RuntimeError(f"Failed to create graph: {r.content}")
def post(uri, data, headers): if getattr(ArgumentParser.args, "psk", None): encode_jwt_to_headers(headers, {}, ArgumentParser.args.psk) r = requests.post(uri, data=data, headers=headers, stream=True) if r.status_code != 200: log.error(r.content.decode()) raise RuntimeError(f"Failed to query graph: {r.content.decode()}") for line in r.iter_lines(): if not line: continue try: data = json.loads(line.decode("utf-8")) yield data except TypeError as e: log.error(e) continue
def registration( self, action: str, client: Callable, data: Optional[Dict] = None ) -> bool: url = f"{self.resotocore_uri}/subscriber/{self.identifier}/{action}" headers = {"accept": "application/json"} if getattr(ArgumentParser.args, "psk", None): encode_jwt_to_headers(headers, {}, ArgumentParser.args.psk) r = client(url, headers=headers, params=data) if r.status_code != 200: raise RuntimeError( f'Error during (un)registration for "{action}"' f" actions: {r.content.decode('utf-8')}" ) return True
def connect(self) -> None: for event, data in self.actions.items(): if not isinstance(data, dict): data = None self.register(event, data) ws_uri = f"{self.resotocore_ws_uri}/subscriber/{self.identifier}/handle" log.debug(f"{self.identifier} connecting to {ws_uri}") headers = {} if getattr(ArgumentParser.args, "psk", None): encode_jwt_to_headers(headers, {}, ArgumentParser.args.psk) self.ws = websocket.WebSocketApp( ws_uri, header=headers, on_open=self.on_open, on_message=self.on_message, on_error=self.on_error, on_close=self.on_close, ) self.ws.run_forever()
def query(query_str: str, query_uri: str) -> Iterator: headers = {"Accept": "application/x-ndjson"} if ArgumentParser.args.psk: encode_jwt_to_headers(headers, {}, ArgumentParser.args.psk) r = requests.post( query_uri, data=query_str, headers=headers, stream=True, ) if r.status_code != 200: raise RuntimeError(f"Failed to query graph: {r.content.decode('utf-8')}") for line in r.iter_lines(): if not line: continue data = json.loads(line.decode("utf-8")) yield data
def search(search_str: str, search_uri: str, tls_data: Optional[TLSData] = None) -> Iterator: headers = {"Accept": "application/x-ndjson"} if ArgumentParser.args.psk: encode_jwt_to_headers(headers, {}, ArgumentParser.args.psk) r = requests.post( search_uri, data=search_str, headers=headers, stream=True, verify=getattr(tls_data, "verify", None), ) if r.status_code != 200: raise RuntimeError(f"Failed to search graph: {r.content.decode('utf-8')}") for line in r.iter_lines(): if not line: continue data = json.loads(line.decode("utf-8")) yield data
def connect(self) -> None: log.debug(f"Connecting to {self.ws_uri}") headers = {} if getattr(ArgumentParser.args, "psk", None): encode_jwt_to_headers(headers, {}, ArgumentParser.args.psk) self.ws = websocket.WebSocketApp( self.ws_uri, header=headers, on_open=self.on_open, on_message=self.on_message, on_error=self.on_error, on_close=self.on_close, on_ping=self.on_ping, on_pong=self.on_pong, ) sslopt = None if self.tls_data: sslopt = {"ca_certs": self.tls_data.ca_cert_path} self.ws.run_forever(sslopt=sslopt, ping_interval=30, ping_timeout=10, ping_payload="ping")
def connect(self) -> None: resotocore_ws_uri_split = urlsplit(self.resotocore_ws_uri) scheme = resotocore_ws_uri_split.scheme netloc = resotocore_ws_uri_split.netloc path = "/work/queue" query_dict = {"task": ",".join(self.tasks)} query_dict.update({k: ",".join(v) for k, v in self.task_queue_filter.items()}) query = urlencode(query_dict) ws_uri = urlunsplit((scheme, netloc, path, query, "")) log.debug(f"{self.identifier} connecting to {ws_uri}") headers = {} if getattr(ArgumentParser.args, "psk", None): encode_jwt_to_headers(headers, {}, ArgumentParser.args.psk) self.ws = websocket.WebSocketApp( ws_uri, header=headers, on_open=self.on_open, on_message=self.on_message, on_error=self.on_error, on_close=self.on_close, ) self.ws.run_forever()
def test_jwt(): psk = "somesecret" payload = {"Hello": "World"} j1 = encode_jwt(payload, psk, expire_in=-1) j2 = encode_jwt(payload, psk, expire_in=-1) past_date = datetime.datetime.now() - datetime.timedelta(30) future_date = datetime.datetime.now() + datetime.timedelta(30) expired_payload = {"Hello": "World", "exp": past_date} future_payload = {"Hello": "World", "exp": future_date} nbf_payload = {"Hello": "World", "nbf": future_date} expired_jwt = encode_jwt(expired_payload, psk) valid_jwt = encode_jwt(future_payload, psk) not_yet_valid_jwt = encode_jwt(nbf_payload, psk) h1 = encode_jwt_to_headers({}, payload, psk, expire_in=-1) assert j1 != j2 assert decode_jwt(j1, psk) == payload assert decode_jwt(j2, psk) == payload assert decode_jwt_from_headers(h1, psk) == payload assert decode_jwt_from_headers({}, psk) is None assert decode_jwt_from_header_value("", psk) is None with pytest.raises(InvalidSignatureError) as e: decode_jwt(j1, "wrongpsk") assert str(e.value) == "Signature verification failed" assert decode_jwt(expired_jwt, psk, options={ "verify_exp": False }).get("Hello") == "World" with pytest.raises(ExpiredSignatureError) as e: decode_jwt(expired_jwt, psk) assert str(e.value) == "Signature has expired" assert decode_jwt(valid_jwt, psk).get("Hello") == "World" with pytest.raises(ImmatureSignatureError) as e: decode_jwt(not_yet_valid_jwt, psk) assert str(e.value).startswith("The token is not yet valid") with pytest.raises(MissingRequiredClaimError) as e: decode_jwt(j1, psk, options={"require": ["exp"]}) assert str(e.value) == 'Token is missing the "exp" claim'