def test_FrozenDict(self):
d = FrozenDict(cmd_args)
with self.subTest():
with self.assertRaises(TypeError):
d.origin = "ORIGIN"
with self.subTest():
self.assertEqual(d.content.action, "deny")
def _subscribe(self, data: FrozenDict) -> None:
close = set()
for args in data.values():
close.add("{host}:{port}".format(**args))
self._check_subscribe(args)
if diff := {*self._clients.keys()} - close:
self.shutdown(list(diff))
def _socketMsg(self, action, act_args, *args, **kwargs):
auth = kwargs.get("headers", {}).get("Authorization", "")
token = re.sub(r"^JWT\s+", "", auth) if auth.startswith("JWT") else ""
url = f"api{act_args['url'].format(*args)}"
url_params = act_args.get('params', {})
if len(url_params) > 0:
url += f"?{'&'.join(f'{k}={v}' for k, v in url_params.items())}"
rtn = dict(
body={},
method=act_args["method"],
status_code=500,
url=f"{self._root_url}{url}",
# Extra Options
meta={})
try:
self._webSocket.send(
json.dumps(
dict(endpoint=url,
method=act_args["method"],
jwt=token,
data=kwargs.get("body", {}),
types=dict(
success=f"@@socket/{action.upper()}_SUCCESS",
failure=f"@@socket/{action.upper()}_FAILURE"))))
try:
rslt = json.loads(self._webSocket.recv())
except ValueError as e:
rslt = {}
rtn.update(
body=rslt.get('payload', {}),
status_code=safe_cast(
rslt.get('meta', {}).get("status_code", 200), int, 200),
# Extra Options
meta=rslt.get('meta', {}))
return FrozenDict(rtn)
except Exception as e:
print(e)
rtn.update(status_code=500, )
return FrozenDict(rtn)
def __init__(self, root: str = _ROOT_DIR, act_id: str = _ACT_ID) -> None:
"""
Initialize and start the Actuator Process
:param root: rood directory of actuator - default CWD
:param act_id: id of the actuator - default UUIDv4
"""
config_file = os.path.join(root, "config.json")
schema_file = os.path.join(root, "schema.json")
config = general.safe_load(config_file)
if len({"actuator_id", "schema"} - set(config.keys())) != 0:
config.setdefault("actuator_id", act_id)
config.setdefault("schema", general.safe_load(schema_file))
json.dump(config, open(config_file, "w"), indent=4)
self._config = FrozenDict(config)
self._dispatch = dispatch.Dispatch(
act=self, dispatch_transform=self._dispatch_transform)
self._dispatch.register(exceptions.action_not_implemented, "default")
self._pairs = None
self._valid_actions = ()
self._valid_targets = ()
# Get valid Actions & Targets from the schema
if len({"meta", "types"} -
set(self._config.schema.keys())) == 0: # JADN
self._profile = self._config.schema.get("meta", {}).get(
"title", "N/A").replace(" ", "_").lower()
self._validator = None
for key in ("Action", "Target"):
key_def = [
x for x in self._config.schema.get("types", [])
if x[0] == key
]
key_def = key_def[0] if len(key_def) == 1 else None
if key_def:
setattr(self, f"_valid_{key.lower()}s",
tuple(a[1] for a in key_def[4]))
else:
raise KeyError(f"{key} not found in schema")
else: # JSON
self._profile = self._config.schema.get("title", "N/A").replace(
" ", "_").lower()
self._validator = general.ValidatorJSON(self._config.schema)
schema_defs = self._config.schema.get("definitions", {})
self._valid_actions = tuple(
a["const"]
for a in schema_defs.get("Action", {}).get("oneOf", []))
self._valid_targets = tuple(
schema_defs.get("Target", {}).get("properties", {}).keys())
def _socket(self):
self._webSocket = websocket.create_connection(self._socket_url,
timeout=2)
init_msg = self._webSocket.recv()
api = dict()
for name, cls in self._api.items():
res = {}
for act, args in getattr(cls, "actions", {}).items():
res[act] = partial(self._socketMsg, act, args)
api[name] = FrozenDict(res)
return api
def _rest(self):
api = API(
api_root_url=self._root_url, # base api url
headers={ # default headers
"Content-Type": "application/json"
},
timeout=2, # default timeout in seconds
append_slash=True, # append slash to final url
json_encode_body=True, # encode body as json
)
for name, cls in self._api.items():
api.add_resource(resource_name=name, resource_class=cls)
_api = {}
for resource in api.get_resource_list():
res = getattr(api, resource)
_api[resource] = FrozenDict(
{act: getattr(res, act)
for act in res.actions})
return FrozenDict(_api)
def pairs(self) -> FrozenDict:
"""
Valid Action/Target pairs registered to this actuator instance
:return: Action/Target Pairs
"""
if self._pairs is None:
pairs = {}
for p in self._dispatch.registered:
p = p.split(".")
if "default" not in p:
pairs.setdefault(p[0], []).append(p[1])
self._pairs = FrozenDict(pairs)
return self._pairs
def mqtt_publish(recipients: List[str], source: dict, device: dict,
body: Union[bytes, str], topic: str, auth: Auth,
headers: dict) -> None:
# pylint: disable=unbalanced-tuple-unpacking
(orc_id, corr_id) = destructure(source, "orchestratorID", "correlationID")
# pylint: disable=unbalanced-tuple-unpacking
(fmt, encoding, broker_socket) = destructure(device,
("format", "broadcast"),
("encoding", "json"),
("socket", "localhost:1883"))
(host, port) = broker_socket.split(":", 1)
payload = Message(recipients=recipients,
origin=f"{orc_id}@{broker_socket}",
msg_type=MessageType.Request,
request_id=uuid.UUID(corr_id),
content_type=SerialFormats(encoding)
if encoding in SerialFormats else SerialFormats.JSON,
content=json.loads(body))
print(f"Sending {broker_socket} topic: {topic} -> {payload}")
publish_props = Properties(PacketTypes.PUBLISH)
publish_props.PayloadFormatIndicator = int(
SerialFormats.is_binary(payload.content_type) is False)
publish_props.ContentType = "application/openc2" # Content-Type
publish_props.UserProperty = ("msgType", payload.msg_type) # User Property
publish_props.UserProperty = ("encoding", payload.content_type
) # User Property
try:
publish_single(config=FrozenDict(MQTT_HOST=host,
MQTT_PORT=safe_cast(port, int, 1883),
USERNAME=auth.username,
PASSWORD=auth.password,
TLS_SELF_SIGNED=safe_cast(
os.environ.get(
"MQTT_TLS_SELF_SIGNED", 0),
int, 0),
CAFILE=auth.caCert,
CLIENT_CERT=auth.clientCert,
CLIENT_KEY=auth.clientKey),
topic=topic,
payload=payload.serialize(),
properties=publish_props)
print(f"Placed payload onto topic {topic} Payload Sent: {payload}")
except Exception as e:
print(
f"There was an error sending command to {broker_socket} topic: {topic} -> {e}"
)
send_error_response(e, headers)
def __init__(self, root: str = _ROOT_DIR, act_id: str = _ACT_ID) -> None:
"""
Initialize and start the Actuator Process
:param root: rood directory of actuator - default CWD
:param act_id: id of the actuator - default UUIDv4
"""
config_file = os.path.join(root, "config.json")
schema_file = os.path.join(root, "schema.json")
config = general.safe_load(config_file)
if "actuator_id" not in config.keys():
config.setdefault("actuator_id", act_id)
json.dump(config, open(config_file, "w"), indent=4)
# Initialize etcd client
self.etcdClient = etcd.Client(host=os.environ.get('ETCD_HOST', 'etcd'),
port=safe_cast(
os.environ.get('ETCD_PORT', 4001),
int, 4001))
schema = general.safe_load(schema_file)
self._config = FrozenDict(**config, schema=schema)
self._dispatch = dispatch.Dispatch(
act=self, dispatch_transform=self._dispatch_transform)
self._dispatch.register(exceptions.action_not_implemented, "default")
self._pairs = None
# Get valid Actions & Targets from the schema
self._profile = self._config.schema.get("title",
"N/A").replace(" ",
"_").lower()
self._validator = general.ValidatorJSON(schema)
schema_defs = self._config.schema.get("definitions", {})
self._prefix = '/actuator'
profiles = self.nsid if len(self.nsid) > 0 else [self._profile]
for profile in profiles:
self.etcdClient.write(f"{self._prefix}/{profile}",
self._config.actuator_id)
self._valid_actions = tuple(
a["const"] for a in schema_defs.get("Action", {}).get("oneOf", []))
self._valid_targets = tuple(
schema_defs.get("Target", {}).get("properties", {}).keys())
def __init__(self, root=ROOT_DIR, act_id=ACT_ID, enable_etcd=True) -> None:
"""
Initialize and start the Actuator Process
:param root: rood directory of actuator - default CWD
:param act_id: id of the actuator - default UUIDv4
"""
config_file = os.path.join(root, "config.json")
schema_file = os.path.join(root, "schema.json")
# Set config
config = general.safe_load(config_file)
if "actuator_id" not in config.keys():
config.setdefault("actuator_id", act_id)
with open(config_file, "w", encoding="UTF-8") as f:
json.dump(config, f, indent=4)
schema = general.safe_load(schema_file)
self._config = FrozenDict(
**config,
schema=schema
)
# Configure Action/Target functions
self._dispatch = dispatch.Dispatch(namespace="root", dispatch_transform=self._dispatch_transform, act=self)
self._dispatch.register(exceptions.action_not_implemented, "default")
# Get valid Actions & Targets from the schema
self._validator = general.ValidatorJSON(schema)
self._profile = self._config.schema.get("title", "N/A").replace(" ", "_").lower()
schema_defs = self._config.schema.get("definitions", {})
self._valid_actions = tuple(a["const"] for a in schema_defs.get("Action", {}).get("oneOf", []))
self._valid_targets = tuple(schema_defs.get("Target", {}).get("properties", {}).keys())
# Initialize etcd client and set profiles
if enable_etcd:
self._etcd = etcd.Client(
host=os.environ.get('ETCD_HOST', 'etcd'),
port=safe_cast(os.environ.get('ETCD_PORT', 4001), int, 4001)
)
profiles = self.nsid if len(self.nsid) > 0 else [self._profile]
for profile in profiles:
self._etcd.write(f"{self._prefix}/{profile}", self._config.actuator_id)
def update(self, data: FrozenDict) -> None:
for t_id, args in data.items():
self._check_subscribe(args)
class MessageQueue:
_auth = FrozenDict({'username': '******', 'password': '******'})
_exchange = 'orchestrator'
_consumerKey = 'response'
_producerExchange = 'producer_transport'
def __init__(self,
hostname='127.0.0.1',
port=5672,
auth=_auth,
exchange=_exchange,
consumer_key=_consumerKey,
producer_exchange=_producerExchange,
callbacks=None):
"""
Message Queue - holds a consumer class and producer class for ease of use
:param hostname: server ip/hostname to connect
:param port: port the AMQP Queue is listening
:param exchange: name of the default exchange
:param consumer_key: key to consumer
:param producer_exchange: ...
:param callbacks: list of functions to call on message receive
"""
self._exchange = exchange if isinstance(exchange,
str) else self._exchange
self._consumerKey = consumer_key if isinstance(
consumer_key, str) else self._consumerKey
self._producerExchange = producer_exchange if isinstance(
producer_exchange, str) else self._producerExchange
self._publish_opts = dict(host=hostname, port=safe_cast(port, int))
self._consume_opts = dict(host=hostname,
port=safe_cast(port, int),
exchange=self._exchange,
routing_key=self._consumerKey,
callbacks=callbacks)
self.producer = Producer(**self._publish_opts)
self.consumer = Consumer(**self._consume_opts)
def send(self, msg, headers, exchange=_producerExchange, routing_key=None):
"""
Publish a message to the specified que and transport
:param msg: message to be published
:param headers: header information for the message being sent
:param exchange: exchange name
:param routing_key: routing key name
:return: None
"""
headers = headers or {}
if routing_key is None:
raise ValueError('Routing Key cannot be None')
self.producer.publish(message=msg,
headers=headers,
exchange=exchange,
routing_key=routing_key)
def shutdown(self):
"""
Shutdown the connection to the queue
"""
self.consumer.shutdown()
self.consumer.join()
import os
from sb_utils import FrozenDict, safe_cast
Config = FrozenDict(
TLS_ENABLED=os.environ.get('MQTT_TLS_ENABLED', False),
TLS_SELF_SIGNED=safe_cast(os.environ.get('MQTT_TLS_SELF_SIGNED', 0), int,
0),
CAFILE=os.environ.get('MQTT_CAFILE', None),
CLIENT_CERT=os.environ.get('MQTT_CLIENT_CERT', None),
CLIENT_KEY=os.environ.get('MQTT_CLIENT_KEY', None),
USERNAME=os.environ.get('MQTT_DEFAULT_USERNAME', None),
PASSWORD=os.environ.get('MQTT_DEFAULT_PASSWORD', None),
MQTT_PREFIX=os.environ.get('MQTT_PREFIX', ''),
MQTT_HOST=os.environ.get('MQTT_HOST', 'queue'),
MQTT_PORT=safe_cast(os.environ.get('MQTT_PORT', 1883), int, 1883),
# TODO: find alternatives??
TRANSPORT_TOPICS=[
t.lower().strip()
for t in os.environ.get("MQTT_TRANSPORT_TOPICS", "").split(",")
],
TOPICS=[
t.lower().strip() for t in os.environ.get("MQTT_TOPICS", "").split(",")
],
# ETCD Options
ETCD_HOST=os.environ.get('ETCD_HOST', 'etcd'),
ETCD_PORT=safe_cast(os.environ.get('ETCD_PORT', 2379), int, 2379))
"""
Query Target functions
"""
from sb_utils import FrozenDict
from sb_utils.actuator import Dispatch, exceptions
Query = Dispatch("query")
Features = FrozenDict(
pairs=lambda act: act.pairs,
profiles=lambda act: act.profile,
rate_limit=lambda act: getattr(act, "rate_limit", 0),
versions=lambda act: act.schema.get("meta", {}).get("version", "N/A"))
@Query.register
def default(*extra_args, **extra_kwargs):
return exceptions.target_not_implemented()
@Query.register
def features(act, target=[], args={}, *extra_args, **extra_kwargs):
if not isinstance(args, dict) and len(set(args) - {"response"}) > 0:
print("Invalid Query Args")
return exceptions.bad_argument()
if not isinstance(target,
list) and len(set(target) - set(Features.keys())) > 0:
return exceptions.bad_request()
else:
from sb_utils import FrozenDict
default_app_config = 'tracking.conf.TrackingConfig'
LEVELS = ('Debug', 'Error', 'Fatal', 'Info', 'Trace', 'Warn')
_DB_LEVELS = tuple((l[0].upper(), l) for l in LEVELS)
EVENT_LEVELS = FrozenDict({l: l[0].upper() for l in LEVELS})
LEVEL_EVENTS = FrozenDict(map(reversed, EVENT_LEVELS.items()))
REQUEST_LEVELS = FrozenDict(Information=range(100, 199),
Success=range(200, 299),
Redirect=range(300, 399),
Client_Error=range(400, 499),
Server_Error=range(500, 599))
