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'
