diff --git a/gremlinapi/attack_helpers.py b/gremlinapi/attack_helpers.py index c74970a..1415e72 100644 --- a/gremlinapi/attack_helpers.py +++ b/gremlinapi/attack_helpers.py @@ -32,7 +32,7 @@ def __init__(self, *args: tuple, **kwargs: dict): self.strategy_type = kwargs.get("strategy_type", "random") # type: ignore def target_definition(self) -> dict: - model: dict = self.repr_model() + model: dict = self.api_model() _target_definition: dict = { "strategyType": self.strategy_type, "strategy": dict(), @@ -60,7 +60,7 @@ def target_definition(self) -> dict: return _target_definition def target_definition_graph(self) -> dict: - model: dict = self.repr_model() + model: dict = self.api_model() _target_definition: dict = { "strategy_type": self.strategy_type, "strategy": dict(), @@ -144,7 +144,7 @@ def strategy_type(self, _target_type: str = None) -> None: _target_type.lower(), None ) - def repr_model(self) -> dict: + def api_model(self) -> dict: model: dict = dict() model["type"] = self.strategy_type if self.exact >= 0 and not self.percent: @@ -158,8 +158,14 @@ def repr_model(self) -> dict: return model def __repr__(self) -> str: - model: dict = self.repr_model() - return json.dumps(model) + kwargs: dict = {} + kwargs["exact"] = self.exact + kwargs["percent"] = self.percent + kwargs["strategy_type"] = self.strategy_type + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) class GremlinAttackCommandHelper(object): @@ -183,7 +189,7 @@ def __init__(self, *args: tuple, **kwargs: dict): self.length = kwargs.get("length", 60) # type: ignore def impact_definition(self) -> dict: - model: dict = self.repr_model() + model: dict = self.api_model() _impact_definition: dict = { "commandArgs": {"cliArgs": [str(self.shortType)], "length": self.length}, "commandType": str(self.shortType), @@ -192,7 +198,7 @@ def impact_definition(self) -> dict: return _impact_definition def impact_definition_graph(self) -> dict: - model: dict = self.repr_model() + model: dict = self.api_model() _impact_definition: dict = { "infra_command_args": { "cli_args": [str(self.shortType)], @@ -260,7 +266,7 @@ def shortType(self, _shortType: str = None) -> None: self._commandType = self._typeMap[_shortType] self._shortType = _shortType - def repr_model(self) -> dict: + def api_model(self) -> dict: model: dict = { "type": self.shortType, "commandType": self.commandType, @@ -269,8 +275,12 @@ def repr_model(self) -> dict: return model def __repr__(self) -> str: - model: dict = self.repr_model() - return json.dumps(model) + kwargs: dict = {} + kwargs["length"] = self.length + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) class GremlinAttackHelper(object): @@ -316,16 +326,21 @@ def target(self, _target: GremlinAttackTargetHelper) -> None: raise GremlinCommandTargetError(error_msg) self._target = _target - def repr_model(self) -> dict: + def api_model(self) -> dict: model: dict = { - "target": json.loads(str(self.target)), - "command": json.loads(str(self.command)), + "target": self.target.api_model(), + "command": self.command.api_model(), } return model def __repr__(self) -> str: - model: dict = self.repr_model() - return json.dumps(model) + kwargs: dict = {} + kwargs["target"] = repr(self.target) + kwargs["command"] = repr(self.command) + return "%s(%s)" % (self.__class__.__name__, kwargs) + + def __str__(self) -> str: + return repr(self) class GremlinTargetHosts(GremlinAttackTargetHelper): @@ -450,8 +465,8 @@ def _valid_tag_pair(self, tagKey: str = None, tagValue: str = None) -> bool: return True return False - def repr_model(self) -> dict: - model: dict = super().repr_model() + def api_model(self) -> dict: + model: dict = super().api_model() if self.target_all_hosts: model["hosts"] = "all" else: @@ -461,20 +476,16 @@ def repr_model(self) -> dict: model["hosts"] = {"multiSelectTags": self.tags} return model - # def __repr__(self): - # model = json.loads(super().__repr__()) - # if self.target_all_hosts: - # model['hosts'] = 'all' - # else: - # if len(self.ids) > 0: - # model['hosts'] = { - # 'ids': self.ids - # } - # elif len(self.tags) > 0: - # model['hosts'] = { - # 'multiSelectTags': self.tags - # } - # return json.dumps(model) + def __repr__(self) -> str: + kwargs: dict = {} + kwargs["exact"] = self.exact + kwargs["percent"] = self.percent + kwargs["strategy_type"] = self.strategy_type + kwargs["target_all_hosts"] = self.target_all_hosts + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) class GremlinTargetContainers(GremlinAttackTargetHelper): @@ -617,8 +628,8 @@ def _valid_label_pair(self, labelKey: Any = None, labelValue: Any = None) -> boo return True return False - def repr_model(self) -> dict: - model: dict = super().repr_model() + def api_model(self) -> dict: + model: dict = super().api_model() if self.target_all_containers: model["containers"] = "all" else: @@ -628,20 +639,18 @@ def repr_model(self) -> dict: model["containers"] = {"multiSelectLabels": self.labels} return model - # def __repr__(self): - # model = json.loads(super().__repr__()) - # if self.target_all_containers: - # model['containers'] = 'all' - # else: - # if len(self.ids) > 0: - # model['containers'] = { - # 'ids': self.ids - # } - # elif len(self.labels) > 0: - # model['containers'] = { - # 'multiSelectLabels': self.labels - # } - # return json.dumps(model) + def __repr__(self) -> str: + kwargs: dict = {} + kwargs["exact"] = self.exact + kwargs["percent"] = self.percent + kwargs["strategy_type"] = self.strategy_type + kwargs["target_all_containers"] = self.target_all_containers + kwargs["ids"] = self.ids + kwargs["labels"] = self.labels + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) class GremlinResourceAttackHelper(GremlinAttackCommandHelper): @@ -700,11 +709,35 @@ def workers(self, _workers: int = None) -> None: raise GremlinParameterError(error_msg) self._workers = _workers + def api_model(self) -> dict: + model: dict = super().api_model() + return model + + def __repr__(self) -> str: + kwargs: dict = {} + kwargs["length"] = self.length + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) + class GremlinStateAttackHelper(GremlinAttackCommandHelper): def __init__(self, *args: tuple, **kwargs: dict): super().__init__(*args, **kwargs) + def api_model(self) -> dict: + model: dict = super().api_model() + return model + + def __repr__(self) -> str: + kwargs: dict = {} + kwargs["length"] = self.length + return "%s()" % (self.__class__.__name__) + + def __str__(self) -> str: + return repr(self) + class GremlinNetworkAttackHelper(GremlinAttackCommandHelper): def __init__(self, *args: tuple, **kwargs: dict): @@ -918,8 +951,8 @@ def tags(self, _tags: Union[list, dict] = None) -> None: self._ids = [] self.target_all_hosts = False - def repr_model(self) -> dict: - model: dict = super().repr_model() + def api_model(self) -> dict: + model: dict = super().api_model() if self.device: model["args"].extend(["-d", self.device]) if len(self.ips) > 0: @@ -932,21 +965,18 @@ def repr_model(self) -> dict: model["trafficImpactMapping"] = {"multiSelectTags": self.tags} return model - # def __repr__(self): - # model = json.loads(super().__repr__()) - # if self.device: - # model['args'].extend(['-d', self.device]) - # if len(self.ips) > 0: - # model['args'].extend(['-i', ','.join(self.ips)]) - # if self.protocol: - # model['args'].extend(['-P', self.protocol]) - # if self.providers and len(self.providers) > 0: - # model['providers'] = self.providers - # if self.tags and len(self.tags) > 0: - # model['trafficImpactMapping'] = { - # 'multiSelectTags': self.tags - # } - # return json.dumps(model) + def __repr__(self) -> str: + kwargs: dict = {} + kwargs["length"] = self.length + kwargs["device"] = self.device + kwargs["ips"] = self.ips + kwargs["protocol"] = self.protocol + kwargs["providers"] = self.providers + kwargs["tags"] = self.tags + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) class GremlinCPUAttack(GremlinResourceAttackHelper): @@ -996,8 +1026,8 @@ def cores(self, _cores: int = None) -> None: raise GremlinParameterError(error_msg) self._cores = _cores - def repr_model(self) -> dict: - model: dict = super().repr_model() + def api_model(self) -> dict: + model: dict = super().api_model() model["args"].extend(["-p", str(self.capacity)]) if self.all_cores: model["args"].append("-a") @@ -1005,14 +1035,16 @@ def repr_model(self) -> dict: model["args"].extend(["-c", str(self.cores)]) return model - # def __repr__(self): - # model = json.loads(super().__repr__()) - # model['args'].extend(['-p', str(self.capacity)]) - # if self.all_cores: - # model['args'].append('-a') - # else: - # model['args'].extend(['-c', str(self.cores)]) - # return json.dumps(model) + def __repr__(self) -> str: + kwargs: dict = {} + kwargs["all_cores"] = self.all_cores + kwargs["capacity"] = self.capacity + kwargs["cores"] = self.cores + kwargs["length"] = self.length + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) class GremlinMemoryAttack(GremlinResourceAttackHelper): @@ -1070,8 +1102,8 @@ def amountType(self, _amountType: str = None) -> None: raise GremlinParameterError(error_msg) self._amountType = _amountType.upper() - def repr_model(self) -> dict: - model: dict = super().repr_model() + def api_model(self) -> dict: + model: dict = super().api_model() if self.amountType == "MB": model["args"].extend(["-m", str(self.amount)]) elif self.amountType == "GB": @@ -1084,19 +1116,15 @@ def repr_model(self) -> dict: raise GremlinParameterError(error_msg) return model - # def __repr__(self): - # model = json.loads(super().__repr__()) - # if self.amountType == 'MB': - # model['args'].extend(['-m', str(self.amount)]) - # elif self.amountType == 'GB': - # model['args'].extend(['-g', str(self.amount)]) - # elif self.amountType == '%': - # model['args'].extend(['-p', str(self.amount)]) - # else: - # error_msg = f'Fatal error, repr_model model may be corrupted, amountType: {self._amountType} is not valid' - # log.fatal(error_msg) - # raise GremlinParameterError(error_msg) - # return json.dumps(model) + def __repr__(self) -> str: + kwargs: dict = {} + kwargs["length"] = self.length + kwargs["amount"] = self.amount + kwargs["amountType"] = self.amountType + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) class GremlinDiskSpaceAttack(GremlinResourceAttackHelper): @@ -1108,21 +1136,25 @@ def __init__(self, *args: tuple, **kwargs: dict): self.percent: int = kwargs.get("percent", 100) # type: ignore self.workers: int = kwargs.get("workers", 1) # type: ignore - def repr_model(self) -> dict: - model: dict = super().repr_model() + def api_model(self) -> dict: + model: dict = super().api_model() model["args"].extend(["-d", str(self.directory)]) model["args"].extend(["-w", str(self.workers)]) model["args"].extend(["-b", str(self.blocksize)]) model["args"].extend(["-p", str(self.percent)]) return model - # def __repr__(self): - # model = json.loads(super().__repr__()) - # model['args'].expand(['-d', str(self.directory)]) - # model['args'].expand(['-w', str(self.workers)]) - # model['args'].expand(['-b', str(self.blocksize)]) - # model['args'].expand(['-p', str(self.percent)]) - # return json.dumps(model) + def __repr__(self) -> str: + kwargs: dict = {} + kwargs["length"] = self.length + kwargs["blocksize"] = self.blocksize + kwargs["directory"] = self.directory + kwargs["percent"] = self.percent + kwargs["workers"] = self.workers + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) class GremlinDiskIOAttack(GremlinResourceAttackHelper): @@ -1163,8 +1195,8 @@ def mode(self, _mode: str = None) -> None: raise GremlinParameterError(error_msg) self._mode = _mode.lower() - def repr_model(self) -> dict: - model: dict = super().repr_model() + def api_model(self) -> dict: + model: dict = super().api_model() model["args"].extend(["-c", str(self.blockcount)]) model["args"].extend(["-d", self.directory]) model["args"].extend(["-m", self.mode]) @@ -1172,14 +1204,18 @@ def repr_model(self) -> dict: model["args"].extend(["-w", str(self.workers)]) return model - # def __repr__(self): - # model = json.loads(super().__repr__()) - # model['args'].extend(['-c', str(self.blockcount)]) - # model['args'].extend(['-d', self.directory]) - # model['args'].extend(['-m', self.mode]) - # model['args'].extend(['-s', str(self.blocksize)]) - # model['args'].extend(['-w', str(self.workers)]) - # return json.dumps(model) + def __repr__(self) -> str: + kwargs: dict = {} + kwargs["length"] = self.length + kwargs["blockcount"] = self.blockcount + kwargs["blocksize"] = self.blocksize + kwargs["directory"] = self.directory + kwargs["mode"] = self.mode + kwargs["workers"] = self.workers + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) class GremlinShutdownAttack(GremlinStateAttackHelper): @@ -1215,19 +1251,22 @@ def reboot(self, _reboot: bool = None) -> None: raise GremlinParameterError(error_msg) self._reboot = _reboot - def repr_model(self) -> dict: - model: dict = super().repr_model() + def api_model(self) -> dict: + model: dict = super().api_model() model["args"] = ["-d", str(self.delay)] if self.reboot: model["args"].append("-r") return model - # def __repr__(self): - # model = json.loads(super().__repr__()) - # model['args'] = ['-d', str(self.delay)] - # if self.reboot: - # model['args'].append('-r') - # return json.dumps(model) + def __repr__(self) -> str: + kwargs: dict = {} + kwargs["length"] = self.length + kwargs["delay"] = self.delay + kwargs["reboot"] = self.reboot + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) class GremlinProcessKillerAttack(GremlinStateAttackHelper): @@ -1376,8 +1415,8 @@ def user(self, _user: str = None) -> None: log.error(error_msg) raise GremlinParameterError(error_msg) - def repr_model(self) -> dict: - model: dict = super().repr_model() + def api_model(self) -> dict: + model: dict = super().api_model() model["args"].extend(["-i", str(self.interval)]) if self.group: model["args"].extend(["-g", self.group]) @@ -1401,30 +1440,22 @@ def repr_model(self) -> dict: model["args"].append("-o") return model - # def __repr__(self): - # model = json.loads(super().__repr__()) - # model['args'].extend(['-i', str(self.interval)]) - # if self.group: - # model['args'].extend(['-g', self.group]) - # if self.process: - # model['args'].extend(['-p', self.process]) - # else: - # error_msg = f'process is required to a be a non-empty string' - # log.fatal(error_msg) - # raise GremlinParameterError(error_msg) - # if self.user: - # model['args'].extend(['-u', self.user]) - # if self.exact: - # model['args'].append('-e') - # if self.full_match: - # model['args'].append('-f') - # if self.kill_children: - # model['args'].append('-c') - # if self.target_newest and not self.target_oldest: - # model['args'].append('-n') - # if self.target_oldest and not self.target_newest: - # model['args'].append('-o') - # return json.dumps(model) + def __repr__(self) -> str: + kwargs: dict = {} + kwargs["length"] = self.length + kwargs["exact"] = self.exact + kwargs["full_match"] = self.full_match + kwargs["group"] = self.group + kwargs["interval"] = self.interval + kwargs["kill_children"] = self.kill_children + kwargs["process"] = self.process + kwargs["target_newest"] = self.target_newest + kwargs["target_oldest"] = self.target_oldest + kwargs["user"] = self.user + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) class GremlinTimeTravelAttack(GremlinStateAttackHelper): @@ -1460,19 +1491,22 @@ def offset(self, offset: int = None) -> None: raise GremlinParameterError(error_msg) self._offset = offset - def repr_model(self) -> dict: - model: dict = super().repr_model() + def api_model(self) -> dict: + model: dict = super().api_model() model["args"].extend(["-o", str(self.offset)]) if self.block_ntp: model["args"].append("-n") return model - # def __repr__(self): - # model = json.loads(super().__repr__()) - # model['args'].extend(['-o', str(self.offset)]) - # if self.block_ntp: - # model['args'].append('-n') - # return json.dumps(model) + def __repr__(self) -> str: + kwargs: dict = {} + kwargs["length"] = self.length + kwargs["block_ntp"] = self.block_ntp + kwargs["offset"] = self.offset + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) class GremlinBlackholeAttack(GremlinNetworkAttackHelper): @@ -1483,8 +1517,8 @@ def __init__(self, *args: tuple, **kwargs: dict): self.hostnames: str = kwargs.get("hostnames", "^api.gremlin.com") # type: ignore self.ingress_ports: list = kwargs.get("ingress_ports", []) # type: ignore - def repr_model(self) -> dict: - model: dict = super().repr_model() + def api_model(self) -> dict: + model: dict = super().api_model() if len(self.egress_ports) > 0: model["args"].extend(["-p", ",".join(self.egress_ports)]) if len(self.hostnames) > 0: @@ -1494,14 +1528,20 @@ def repr_model(self) -> dict: return model def __repr__(self) -> str: - model: dict = json.loads(super().__repr__()) - if len(self.egress_ports) > 0: - model["args"].extend(["-p", ",".join(self.egress_ports)]) - if len(self.hostnames) > 0: - model["args"].extend(["-h", ",".join(self.hostnames)]) - if len(self.ingress_ports) > 0: - model["args"].extend(["-n", ",".join(self.ingress_ports)]) - return json.dumps(model) + kwargs: dict = {} + kwargs["length"] = self.length + kwargs["device"] = self.device + kwargs["ips"] = self.ips + kwargs["protocol"] = self.protocol + kwargs["providers"] = self.providers + kwargs["tags"] = self.tags + kwargs["egress_ports"] = self.egress_ports + kwargs["hostnames"] = self.hostnames + kwargs["ingress_ports"] = self.ingress_ports + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) class GremlinDNSAttack(GremlinNetworkAttackHelper): @@ -1511,13 +1551,22 @@ def __init__(self, *args: tuple, **kwargs: dict): self._allowed_protocols: list = ["TCP", "UDP"] self.protocol: str = kwargs.get("protocol", "") # type: ignore - def repr_model(self) -> dict: - model: dict = super().repr_model() + def api_model(self) -> dict: + model: dict = super().api_model() return model - # def __repr__(self): - # model = json.loads(super().__repr__()) - # return json.dumps(model) + def __repr__(self) -> str: + kwargs: dict = {} + kwargs["length"] = self.length + kwargs["device"] = self.device + kwargs["ips"] = self.ips + kwargs["protocol"] = self.protocol + kwargs["providers"] = self.providers + kwargs["tags"] = self.tags + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) class GremlinLatencyAttack(GremlinNetworkAttackHelper): @@ -1542,8 +1591,8 @@ def delay(self, _delay: int = None) -> None: raise GremlinParameterError(error_msg) self._delay = _delay - def repr_model(self) -> dict: - model: dict = super().repr_model() + def api_model(self) -> dict: + model: dict = super().api_model() model["args"].extend(["-m", str(self.delay)]) if len(self.egress_ports) > 0: model["args"].extend(["-p", ",".join(self.egress_ports)]) @@ -1553,16 +1602,22 @@ def repr_model(self) -> dict: model["args"].extend(["-s", ",".join(self.source_ports)]) return model - # def __repr__(self): - # model = json.loads(super().__repr__()) - # model['args'].extend(['-m', str(self.delay)]) - # if len(self.egress_ports) > 0: - # model['args'].extend(['-p', ','.join(self.egress_ports)]) - # if len(self.hostnames) > 0: - # model['args'].extend(['-h', ','.join(self.hostnames)]) - # if len(self.source_ports) > 0: - # model['args'].extend(['-s', ','.join(self.source_ports)]) - # return json.dumps(model) + def __repr__(self) -> str: + kwargs: dict = {} + kwargs["length"] = self.length + kwargs["device"] = self.device + kwargs["ips"] = self.ips + kwargs["protocol"] = self.protocol + kwargs["providers"] = self.providers + kwargs["tags"] = self.tags + kwargs["delay"] = self.delay + kwargs["egress_ports"] = self.egress_ports + kwargs["hostnames"] = self.hostnames + kwargs["source_ports"] = self.source_ports + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) class GremlinPacketLossAttack(GremlinNetworkAttackHelper): @@ -1606,8 +1661,8 @@ def percent(self, _percent: int = None) -> None: raise GremlinParameterError(error_msg) self._percent = _percent - def repr_model(self) -> dict: - model: dict = super().repr_model() + def api_model(self) -> dict: + model: dict = super().api_model() model["args"].extend(["-r", str(self.percent)]) if len(self.egress_ports) > 0: model["args"].extend(["-p", ",".join(self.egress_ports)]) @@ -1619,15 +1674,20 @@ def repr_model(self) -> dict: model["args"].append("-c") return model - # def __repr__(self): - # model = json.loads(super().__repr__()) - # model['args'].extend(['-r', str(self.percent)]) - # if len(self.egress_ports) > 0: - # model['args'].extend(['-p', ','.join(self.egress_ports)]) - # if len(self.hostnames) > 0: - # model['args'].extend(['-h', ','.join(self.hostnames)]) - # if len(self.source_ports) > 0: - # model['args'].extend(['-s', ','.join(self.source_ports)]) - # if self.corrupt: - # model['args'].append('-c') - # return json.dumps(model) + def __repr__(self) -> str: + kwargs: dict = {} + kwargs["length"] = self.length + kwargs["device"] = self.device + kwargs["ips"] = self.ips + kwargs["protocol"] = self.protocol + kwargs["providers"] = self.providers + kwargs["tags"] = self.tags + kwargs["corrupt"] = self.corrupt + kwargs["egress_ports"] = self.egress_ports + kwargs["hostnames"] = self.hostnames + kwargs["percent"] = self.percent + kwargs["source_ports"] = self.source_ports + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) diff --git a/gremlinapi/scenario_graph_helpers.py b/gremlinapi/scenario_graph_helpers.py index d37c047..1ff464c 100644 --- a/gremlinapi/scenario_graph_helpers.py +++ b/gremlinapi/scenario_graph_helpers.py @@ -73,7 +73,7 @@ def add_edge(self, _node: "GremlinScenarioNode", _weight: str = None) -> None: @property def data(self) -> dict: - return self.repr_model() + return self.api_model() @property def id(self) -> str: @@ -151,7 +151,7 @@ def uuid(self) -> str: raise GremlinParameterError(error_msg) return f"{self.name}-{self.id}" - def repr_model(self) -> dict: + def api_model(self) -> dict: model = { "type": self.node_type, "guid": self.uuid, @@ -161,7 +161,12 @@ def repr_model(self) -> dict: return model def __repr__(self) -> str: - return json.dumps(self.repr_model) + kwargs: dict = {} + kwargs["name"] = self.name + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) class GremlinScenarioGraphHelper(object): @@ -339,7 +344,7 @@ def name(self, _name: str = None) -> None: raise GremlinParameterError(error_msg) self._name = _name - def repr_model(self) -> dict: + def api_model(self) -> dict: model: dict = { "description": self.description, "hypothesis": self.hypothesis, @@ -356,8 +361,14 @@ def repr_model(self) -> dict: return model def __repr__(self) -> str: - model: dict = self.repr_model() - return json.dumps(model) + kwargs: dict = {} + kwargs["name"] = self.name + kwargs["description"] = self.description + kwargs["hypothesis"] = self.hypothesis + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) class GremlinScenarioSerialNode(GremlinScenarioNode): @@ -368,6 +379,18 @@ def __init( ): super().__init__(*args, **kwargs) + def api_model(self) -> dict: + model: dict = super().api_model() + return model + + def __repr__(self) -> str: + kwargs: dict = {} + kwargs["name"] = self.name + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) + class GremlinScenarioParallelNode(GremlinScenarioNode): def __init( @@ -378,6 +401,18 @@ def __init( super().__init__(*args, **kwargs) raise NotImplementedError("Parallel Scenario Nodes NOT IMPLEMENTED") + def api_model(self) -> dict: + model: dict = super().api_model() + return model + + def __repr__(self) -> str: + kwargs: dict = {} + kwargs["name"] = self.name + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) + class GremlinScenarioAttackNode(GremlinScenarioSerialNode): def __init__( @@ -403,6 +438,18 @@ def attack_type(self, _attack_type: str = None) -> None: raise GremlinParameterError(error_msg) self._attack_type = _attack_type + def api_model(self) -> dict: + model: dict = super().api_model() + return model + + def __repr__(self) -> str: + kwargs: dict = {} + kwargs["name"] = self.name + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) + class GremlinScenarioILFINode(GremlinScenarioSerialNode): def __init__(self, *args: tuple, **kwargs: dict): @@ -412,8 +459,8 @@ def __init__(self, *args: tuple, **kwargs: dict): self._command: GremlinAttackCommandHelper = None # type: ignore self._target: GremlinAttackTargetHelper = None # type: ignore self.node_type: str = "InfraAttack" - self.command: GremlinAttackCommandHelper = kwargs.get("command", self._command) # type: ignore - self.target: GremlinAttackTargetHelper = kwargs.get("target", self._target) # type: ignore + self.command: GremlinAttackCommandHelper = kwargs.get("command", GremlinAttackCommandHelper()) # type: ignore + self.target: GremlinAttackTargetHelper = kwargs.get("target", GremlinAttackTargetHelper()) # type: ignore @property def command(self) -> GremlinAttackCommandHelper: @@ -427,8 +474,8 @@ def command(self, _command: GremlinAttackCommandHelper) -> None: raise GremlinParameterError(error_msg) self._command = _command - def repr_model(self) -> dict: - model: dict = super().repr_model() + def api_model(self) -> dict: + model: dict = super().api_model() model["impact_definition"] = self.command.impact_definition_graph() model["target_definition"] = self.target.target_definition_graph() return model @@ -445,6 +492,16 @@ def target(self, _target: GremlinAttackTargetHelper) -> None: raise GremlinParameterError(error_msg) self._target = _target + def __repr__(self) -> str: + kwargs: dict = {} + kwargs["name"] = self.name + kwargs["command"] = repr(self.command) + kwargs["target"] = repr(self.target) + return "%s(%s)" % (self.__class__.__name__, kwargs) + + def __str__(self) -> str: + return repr(self) + class GremlinScenarioALFINode(GremlinScenarioSerialNode): def __init__( @@ -456,6 +513,18 @@ def __init__( self.attack_type = "ALFI" raise NotImplementedError("ALFI Scenarios NOT IMPLEMENTED") + def api_model(self) -> dict: + model: dict = super().api_model() + return model + + def __repr__(self) -> str: + kwargs: dict = {} + kwargs["name"] = self.name + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) + class GremlinScenarioDelayNode(GremlinScenarioSerialNode): def __init__(self, *args: tuple, **kwargs: dict): @@ -466,8 +535,8 @@ def __init__(self, *args: tuple, **kwargs: dict): self._delay: int = kwargs.get("delay", None) # type: ignore self.node_type: str = "Delay" - def repr_model(self) -> dict: - model: dict = super().repr_model() + def api_model(self) -> dict: + model: dict = super().api_model() model["delay"] = self.delay return model @@ -485,6 +554,15 @@ def delay(self, _duration: int = None) -> None: raise GremlinParameterError(error_msg) self._delay = _duration + def __repr__(self) -> str: + kwargs: dict = {} + kwargs["name"] = self.name + kwargs["delay"] = self.delay + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) + class GremlinScenarioStatusCheckNode(GremlinScenarioSerialNode): def __init__( @@ -513,8 +591,8 @@ def __init__( "evaluation_response_body_evaluation", "" ) # type: ignore - def repr_model(self) -> dict: - model: dict = super().repr_model() + def api_model(self) -> dict: + model: dict = super().api_model() model["endpointConfiguration"] = { "url": self.endpoint_url, "headers": self.endpoint_headers, @@ -601,6 +679,22 @@ def evaluation_response_body_evaluation( raise GremlinParameterError(error_msg) self._evaluation_response_body_evaluation = _evaluation_response_body_evaluation + def __repr__(self) -> str: + kwargs: dict = {} + kwargs["name"] = self.name + kwargs["description"] = self.description + kwargs["endpoint_url"] = self.endpoint_url + kwargs["endpoint_headers"] = self.endpoint_headers + kwargs["evaluation_ok_status_codes"] = self.evaluation_ok_status_codes + kwargs["evaluation_ok_latency_max"] = self.evaluation_ok_latency_max + kwargs[ + "evaluation_response_body_evaluation" + ] = self.evaluation_response_body_evaluation + return "%s(%s)" % (self.__class__.__name__, json.dumps(kwargs)) + + def __str__(self) -> str: + return repr(self) + class _GremlinNodeGraph(object): def __init__(self): @@ -843,3 +937,14 @@ def _validate_type(self, _node: GremlinScenarioNode) -> bool: log.error(error_msg) raise GremlinParameterError(error_msg) return True + + def api_model(self) -> dict: + model: dict = {} + return model + + def __repr__(self) -> str: + kwargs: dict = {} + return "%s()" % (self.__class__.__name__) + + def __str__(self) -> str: + return repr(self) diff --git a/gremlinapi/util.py b/gremlinapi/util.py index 007cc91..b5896c0 100644 --- a/gremlinapi/util.py +++ b/gremlinapi/util.py @@ -7,6 +7,7 @@ log = logging.getLogger("GremlinAPI.client") + _version = "0.14.3" diff --git a/tests/test_attack_helpers.py b/tests/test_attack_helpers.py index 16011a0..5f23c6e 100644 --- a/tests/test_attack_helpers.py +++ b/tests/test_attack_helpers.py @@ -21,6 +21,7 @@ GremlinDNSAttack, GremlinLatencyAttack, GremlinPacketLossAttack, + GremlinStateAttackHelper, ) from .util import mock_data @@ -39,7 +40,7 @@ class TestAttackHelpers(unittest.TestCase): - def test_attack_helper_repr_model(self) -> None: + def test_attack_helper_api_model(self) -> None: # defaults expected_output = { "command": { @@ -50,9 +51,29 @@ def test_attack_helper_repr_model(self) -> None: "target": {"hosts": "all", "percent": 10, "type": "Random"}, } helper = GremlinAttackHelper() - helper_output = helper.repr_model() + helper_output = helper.api_model() self.assertEqual(helper_output, expected_output) + def test_attack_helper_repr_str(self) -> None: + # defaults + expected_output = 'GremlinAttackHelper({\'target\': \'GremlinTargetHosts({"exact": 0, "percent": 25, "strategy_type": "Random", "target_all_hosts": true})\', \'command\': \'GremlinCPUAttack({"all_cores": false, "capacity": 90, "cores": 1, "length": 60})\'})' + kwargs_th = { + "exact": 0, + "percent": 25, + "strategy_type": "Random", + "target_all_hosts": True, + } + kwargs_cpua = {"all_cores": False, "capacity": 90, "cores": 1, "length": 60} + kwargs = { + "target": GremlinTargetHosts(**kwargs_th), + "command": GremlinCPUAttack(**kwargs_cpua), + } + helper = GremlinAttackHelper(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + def test_target_helper_target_definition(self) -> None: helper = GremlinAttackTargetHelper(**attack_helper_params_custom) test_output = helper.target_definition() @@ -93,13 +114,23 @@ def test_target_helper_target_definition_graph(self) -> None: } self.assertEqual(test_output, expected_output) - def test_attack_target_helper_repr_model(self) -> None: + def test_attack_target_helper_api_model(self) -> None: # defaults expected_output = {"percent": 10, "type": "Random"} helper = GremlinAttackTargetHelper() - helper_output = helper.repr_model() + helper_output = helper.api_model() self.assertEqual(helper_output, expected_output) + def test_attack_target_helper_repr_str(self) -> None: + # defaults + expected_output = 'GremlinAttackTargetHelper({"exact": 0, "percent": 15, "strategy_type": "Random"})' + kwargs = {"exact": 0, "percent": 15, "strategy_type": "Random"} + helper = GremlinAttackTargetHelper(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + @patch("requests.get") def test_filter_active_tags_with_decorator(self, mock_get) -> None: expected_output = { @@ -118,13 +149,28 @@ def mock_json(): helper._filter_active_tags() self.assertEqual(helper._active_tags, expected_output) - def test_target_hosts_repr_model(self) -> None: + def test_target_hosts_api_model(self) -> None: # defaults expected_output = {"hosts": "all", "percent": 10, "type": "Random"} helper = GremlinTargetHosts() - helper_output = helper.repr_model() + helper_output = helper.api_model() self.assertEqual(helper_output, expected_output) + def test_target_hosts_repr_str(self) -> None: + # defaults + expected_output = 'GremlinTargetHosts({"exact": 0, "percent": 25, "strategy_type": "Random", "target_all_hosts": true})' + kwargs = { + "exact": 0, + "percent": 25, + "strategy_type": "Random", + "target_all_hosts": True, + } + helper = GremlinTargetHosts(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + @patch("requests.get") def test_filter_active_labels_with_decorator(self, mock_get) -> None: expected_output = {"os_type": ["testwindows"], "os-version": ["test10"]} @@ -141,13 +187,30 @@ def mock_json(): helper._filter_active_labels() self.assertEqual(helper._active_labels, expected_output) - def test_target_containers_repr_model(self) -> None: + def test_target_containers_api_model(self) -> None: # defaults expected_output = {"containers": "all", "percent": 10, "type": "Random"} helper = GremlinTargetContainers() - helper_output = helper.repr_model() + helper_output = helper.api_model() self.assertEqual(helper_output, expected_output) + def test_target_containers_repr_str(self) -> None: + # defaults + expected_output = 'GremlinTargetContainers({"exact": 0, "percent": 35, "strategy_type": "Random", "target_all_containers": false, "ids": [], "labels": {}})' + kwargs = { + "exact": 0, + "percent": 35, + "strategy_type": "Random", + "target_all_containers": False, + "ids": [], + "labels": {}, + } + helper = GremlinTargetContainers(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + def test_impact_definition(self) -> None: expected_output = { "commandArgs": {"cliArgs": ["memory", "-l", "100"], "length": 100}, @@ -175,20 +238,46 @@ def test_impact_definition_graph(self) -> None: helper_output = helper.impact_definition_graph() self.assertEqual(helper_output, expected_output) - def test_attack_command_helper_repr_model(self) -> None: + def test_attack_command_helper_api_model(self) -> None: # defaults expected_output = {"args": ["-l", "60"], "commandType": "", "type": ""} helper = GremlinAttackCommandHelper() - helper_output = helper.repr_model() + helper_output = helper.api_model() self.assertEqual(helper_output, expected_output) - def test_resource_attack_helper_repr_model(self) -> None: + def test_attack_command_helper_repr_str(self) -> None: + expected_output = 'GremlinAttackCommandHelper({"length": 70})' + kwargs = {"length": 70} + helper = GremlinAttackCommandHelper(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + + def test_resource_attack_helper_api_model(self) -> None: # defaults expected_output = {"args": ["-l", "60"], "commandType": "", "type": ""} helper = GremlinResourceAttackHelper() - helper_output = helper.repr_model() + helper_output = helper.api_model() self.assertEqual(helper_output, expected_output) + def test_resource_attack_helper_repr_str(self) -> None: + expected_output = 'GremlinResourceAttackHelper({"length": 70})' + kwargs = {"length": 70} + helper = GremlinResourceAttackHelper(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + + def test_state_attack_helper_repr_str(self) -> None: + expected_output = "GremlinStateAttackHelper()" + helper = GremlinStateAttackHelper() + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + def test__port_maker(self) -> None: expected_output = [] helper = GremlinNetworkAttackHelper() @@ -205,14 +294,30 @@ def test__port_maker(self) -> None: helper_output = helper._port_maker(expected_output) self.assertEqual(expected_output, helper_output) - def test_network_attack_helper_repr_model(self) -> None: + def test_network_attack_helper_api_model(self) -> None: # defaults expected_output = {"args": ["-l", "60"], "commandType": "", "type": ""} helper = GremlinNetworkAttackHelper() - helper_output = helper.repr_model() + helper_output = helper.api_model() self.assertEqual(helper_output, expected_output) - def test_cpu_attack_repr_model(self) -> None: + def test_network_attack_helper_repr_str(self) -> None: + expected_output = 'GremlinNetworkAttackHelper({"length": 65, "device": "", "ips": [], "protocol": "", "providers": [], "tags": []})' + kwargs = { + "length": 65, + "device": "", + "ips": [], + "protocol": "", + "providers": [], + "tags": [], + } + helper = GremlinNetworkAttackHelper(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + + def test_cpu_attack_api_model(self) -> None: # defaults expected_output = { "args": ["-l", "60", "-p", "100", "-c", "1"], @@ -220,10 +325,20 @@ def test_cpu_attack_repr_model(self) -> None: "type": "cpu", } helper = GremlinCPUAttack() - helper_output = helper.repr_model() + helper_output = helper.api_model() self.assertEqual(helper_output, expected_output) - def test_memory_attack_repr_model(self) -> None: + def test_cpu_attack_repr_str(self) -> None: + # defaults + expected_output = 'GremlinCPUAttack({"all_cores": false, "capacity": 90, "cores": 2, "length": 60})' + kwargs = {"all_cores": False, "capacity": 90, "cores": 2, "length": 60} + helper = GremlinCPUAttack(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + + def test_memory_attack_api_model(self) -> None: # defaults expected_output = { "args": ["-l", "60", "-p", "100"], @@ -231,10 +346,25 @@ def test_memory_attack_repr_model(self) -> None: "type": "memory", } helper = GremlinMemoryAttack() - helper_output = helper.repr_model() + helper_output = helper.api_model() self.assertEqual(helper_output, expected_output) - def test_disk_space_attack_repr_model(self) -> None: + def test_memory_attack_repr_str(self) -> None: + expected_output = ( + 'GremlinMemoryAttack({"length": 65, "amount": 99, "amountType": "MB"})' + ) + kwargs = { + "length": 65, + "amount": 99, + "amountType": "MB", + } + helper = GremlinMemoryAttack(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + + def test_disk_space_attack_api_model(self) -> None: # defaults expected_output = { "args": ["-l", "60", "-d", "/tmp", "-w", "1", "-b", "4", "-p", "100"], @@ -242,10 +372,25 @@ def test_disk_space_attack_repr_model(self) -> None: "type": "disk", } helper = GremlinDiskSpaceAttack() - helper_output = helper.repr_model() + helper_output = helper.api_model() self.assertEqual(helper_output, expected_output) - def test_disk_io_repr_model(self) -> None: + def test_disk_space_attack_repr_str(self) -> None: + expected_output = 'GremlinDiskSpaceAttack({"length": 75, "blocksize": 4, "directory": "/tmp", "percent": 95, "workers": 1})' + kwargs = { + "length": 75, + "blocksize": 4, + "directory": "/tmp", + "percent": 95, + "workers": 1, + } + helper = GremlinDiskSpaceAttack(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + + def test_disk_io_api_model(self) -> None: # defaults expected_output = { "args": [ @@ -266,10 +411,26 @@ def test_disk_io_repr_model(self) -> None: "type": "io", } helper = GremlinDiskIOAttack() - helper_output = helper.repr_model() + helper_output = helper.api_model() self.assertEqual(helper_output, expected_output) - def test_shutdown_attack_repr_model(self) -> None: + def test_disk_io_repr_str(self) -> None: + expected_output = 'GremlinDiskIOAttack({"length": 85, "blockcount": 1, "blocksize": 4, "directory": "/tmp/usr", "mode": "rw", "workers": 2})' + kwargs = { + "length": 85, + "blockcount": 1, + "blocksize": 4, + "directory": "/tmp/usr", + "mode": "rw", + "workers": 2, + } + helper = GremlinDiskIOAttack(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + + def test_shutdown_attack_api_model(self) -> None: # defaults expected_output = { "args": ["-d", "1"], @@ -277,10 +438,25 @@ def test_shutdown_attack_repr_model(self) -> None: "type": "shutdown", } helper = GremlinShutdownAttack() - helper_output = helper.repr_model() + helper_output = helper.api_model() self.assertEqual(helper_output, expected_output) - def test_process_killer_attack_repr_model(self) -> None: + def test_shutdown_attack_repr_str(self) -> None: + expected_output = ( + 'GremlinShutdownAttack({"length": 85, "delay": 20, "reboot": false})' + ) + kwargs = { + "length": 85, + "delay": 20, + "reboot": False, + } + helper = GremlinShutdownAttack(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + + def test_process_killer_attack_api_model(self) -> None: test_process = "ls" kwargs = {"process": test_process} # defaults @@ -290,10 +466,30 @@ def test_process_killer_attack_repr_model(self) -> None: "type": "process_killer", } helper = GremlinProcessKillerAttack(**kwargs) - helper_output = helper.repr_model() + helper_output = helper.api_model() self.assertEqual(helper_output, expected_output) - def test_time_travel_attack_repr_model(self) -> None: + def test_process_killer_attack_repr_str(self) -> None: + expected_output = 'GremlinProcessKillerAttack({"length": 65, "exact": false, "full_match": true, "group": "", "interval": 1, "kill_children": false, "process": "", "target_newest": false, "target_oldest": false, "user": "unittestuser"})' + kwargs = { + "length": 65, + "exact": False, + "full_match": True, + "group": "", + "interval": 1, + "kill_children": False, + "process": "", + "target_newest": False, + "target_oldest": False, + "user": "unittestuser", + } + helper = GremlinProcessKillerAttack(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + + def test_time_travel_attack_api_model(self) -> None: # defaults expected_output = { "args": ["-l", "60", "-o", "86400"], @@ -301,10 +497,23 @@ def test_time_travel_attack_repr_model(self) -> None: "type": "time_travel", } helper = GremlinTimeTravelAttack() - helper_output = helper.repr_model() + helper_output = helper.api_model() self.assertEqual(helper_output, expected_output) - def test_black_hole_attack_repr_model(self) -> None: + def test_time_travel_attack_repr_str(self) -> None: + expected_output = 'GremlinTimeTravelAttack({"length": 65, "block_ntp": false, "offset": 12000})' + kwargs = { + "length": 65, + "block_ntp": False, + "offset": 12000, + } + helper = GremlinTimeTravelAttack(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + + def test_black_hole_attack_api_model(self) -> None: # defaults expected_output = { "args": ["-l", "60", "-p", "^53", "-h", "^api.gremlin.com"], @@ -312,17 +521,52 @@ def test_black_hole_attack_repr_model(self) -> None: "type": "blackhole", } helper = GremlinBlackholeAttack() - helper_output = helper.repr_model() + helper_output = helper.api_model() self.assertEqual(helper_output, expected_output) - def test_dns_attack_repr_model(self) -> None: + def test_black_hole_attack_repr_str(self) -> None: + expected_output = 'GremlinBlackholeAttack({"length": 65, "device": "", "ips": [], "protocol": "", "providers": [], "tags": [], "egress_ports": ["^533"], "hostnames": "^api2.gremlin.com", "ingress_ports": []})' + kwargs = { + "length": 65, + "device": "", + "ips": [], + "protocol": "", + "providers": [], + "tags": [], + "egress_ports": ["^533"], + "hostnames": ["^api2.gremlin.com"], + "ingress_ports": [], + } + helper = GremlinBlackholeAttack(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + + def test_dns_attack_api_model(self) -> None: # defaults expected_output = {"args": ["-l", "60"], "commandType": "DNS", "type": "dns"} helper = GremlinDNSAttack() - helper_output = helper.repr_model() + helper_output = helper.api_model() self.assertEqual(helper_output, expected_output) - def test_latency_attack_repr_model(self) -> None: + def test_dns_attack_repr_str(self) -> None: + expected_output = 'GremlinDNSAttack({"length": 75, "device": "", "ips": [], "protocol": "", "providers": [], "tags": []})' + kwargs = { + "length": 75, + "device": "", + "ips": [], + "protocol": "", + "providers": [], + "tags": [], + } + helper = GremlinDNSAttack(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + + def test_latency_attack_api_model(self) -> None: # defaults expected_output = { "args": ["-l", "60", "-m", "100", "-p", "^53", "-h", "^api.gremlin.com"], @@ -330,10 +574,33 @@ def test_latency_attack_repr_model(self) -> None: "type": "latency", } helper = GremlinLatencyAttack() - helper_output = helper.repr_model() + helper_output = helper.api_model() self.assertEqual(helper_output, expected_output) - def test_packet_loss_attack_repr_model(self) -> None: + def test_latency_attack_repr_str(self) -> None: + expected_output = 'GremlinLatencyAttack({"length": 75, "device": "", "ips": [], "protocol": "", "providers": [], "tags": [], "delay": 100, "egress_ports": ["^533"], "hostnames": "^api2.gremlin.com", "source_ports": []})' + kwargs = { + "length": 75, + "device": "", + "ips": [], + "protocol": "", + "providers": [], + "tags": [], + "delay": 100, + "egress_ports": ["^533"], + "hostnames": ["^api2.gremlin.com"], + "source_ports": [], + } + helper = GremlinLatencyAttack(**kwargs) + helper_output = repr(helper) + max_diff = self.maxDiff = None + self.maxDiff = None + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + self.maxDiff = max_diff + + def test_packet_loss_attack_api_model(self) -> None: # defaults expected_output = { "args": ["-l", "60", "-r", "1", "-p", "^53", "-h", "^api.gremlin.com"], @@ -341,5 +608,29 @@ def test_packet_loss_attack_repr_model(self) -> None: "type": "packet_loss", } helper = GremlinPacketLossAttack() - helper_output = helper.repr_model() + helper_output = helper.api_model() self.assertEqual(helper_output, expected_output) + + def test_packet_loss_attack_repr_str(self) -> None: + expected_output = 'GremlinPacketLossAttack({"length": 85, "device": "", "ips": [], "protocol": "", "providers": [], "tags": [], "corrupt": false, "egress_ports": ["^543"], "hostnames": "^api3.gremlin.com", "percent": 1, "source_ports": []})' + kwargs = { + "length": 85, + "device": "", + "ips": [], + "protocol": "", + "providers": [], + "tags": [], + "corrupt": False, + "egress_ports": ["^543"], + "hostnames": ["^api3.gremlin.com"], + "percent": 1, + "source_ports": [], + } + helper = GremlinPacketLossAttack(**kwargs) + helper_output = repr(helper) + max_diff = self.maxDiff = None + self.maxDiff = None + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + self.maxDiff = max_diff diff --git a/tests/test_scenario_graph_helpers.py b/tests/test_scenario_graph_helpers.py index 7d42d12..cadbc48 100644 --- a/tests/test_scenario_graph_helpers.py +++ b/tests/test_scenario_graph_helpers.py @@ -4,12 +4,21 @@ from gremlinapi.scenario_graph_helpers import ( GremlinScenarioGraphHelper, GremlinScenarioNode, + GremlinScenarioSerialNode, + GremlinScenarioParallelNode, + GremlinScenarioAttackNode, GremlinScenarioILFINode, + GremlinScenarioALFINode, GremlinScenarioDelayNode, GremlinScenarioStatusCheckNode, _GremlinNodeGraph, ) +from gremlinapi.attack_helpers import ( + GremlinAttackTargetHelper, + GremlinAttackCommandHelper, +) + from .util import ( mock_scenario, mock_ilfi_node, @@ -28,7 +37,7 @@ def test_add_node(self) -> None: self.assertEqual(helper._nodes.get_node(helper_node.id), helper_node) self.assertNotEqual(helper._nodes.get_node(helper_node.id), helper_node_2) - def test_gremlin_scenario_graph_helper_repr_model(self) -> None: + def test_gremlin_scenario_graph_helper_api_model(self) -> None: status_check_description = "Check if Gremlin.com is Still Up" endpoint_url = "https://www.google.com" endpoint_headers = dict() @@ -143,10 +152,23 @@ def test_gremlin_scenario_graph_helper_repr_model(self) -> None: "hypothesis": "No Hypothesis", "name": "code_created_scenario_6", } - self.assertEqual(my_scenario.repr_model(), expected_output) + self.assertEqual(my_scenario.api_model(), expected_output) self.maxDiff = t_diff - def test_gremlin_scenario_graph_helper_repr_model_default_add_node(self) -> None: + def test_gremlin_scenario_graph_helper_repr_str(self) -> None: + expected_output = 'GremlinScenarioGraphHelper({"name": "code_created_scenario_6", "description": "Three nodes now", "hypothesis": "No Hypothesis"})' + kwargs = { + "name": "code_created_scenario_6", + "description": "Three nodes now", + "hypothesis": "No Hypothesis", + } + helper = GremlinScenarioGraphHelper(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + + def test_gremlin_scenario_graph_helper_api_model_default_add_node(self) -> None: status_check_description = "Check if Gremlin.com is Still Up" endpoint_url = "https://www.google.com" endpoint_headers = dict() @@ -258,7 +280,7 @@ def test_gremlin_scenario_graph_helper_repr_model_default_add_node(self) -> None "hypothesis": "No Hypothesis", "name": "code_created_scenario_6", } - self.assertEqual(my_scenario.repr_model(), expected_output) + self.assertEqual(my_scenario.api_model(), expected_output) self.maxDiff = t_diff def test_add_edge(self) -> None: @@ -272,7 +294,7 @@ def test_add_edge(self) -> None: self.assertEqual(len(helper._edges), 1) self.assertEqual(helper._edges[helper_2.id]["node"], helper_2) - def test_gremlin_scenario_node_repr_model(self) -> None: + def test_gremlin_scenario_node_api_model(self) -> None: helper = GremlinScenarioNode(**mock_scenario) expected_output = { "guid": "mock_scenario-%s" % helper.id, @@ -281,7 +303,53 @@ def test_gremlin_scenario_node_repr_model(self) -> None: "type": helper.node_type, } - self.assertEqual(helper.repr_model(), expected_output) + self.assertEqual(helper.api_model(), expected_output) + + def test_gremlin_scenario_node_repr_str(self) -> None: + expected_output = 'GremlinScenarioNode({"name": "a mock node"})' + kwargs = { + "name": "a mock node", + } + helper = GremlinScenarioNode(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + + def test_gremlin_scenario_serial_node_repr_str(self) -> None: + expected_output = 'GremlinScenarioSerialNode({"name": "a mock serial node"})' + kwargs = { + "name": "a mock serial node", + } + helper = GremlinScenarioSerialNode(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + + def test_gremlin_scenario_parallel_node_repr_str(self) -> None: + expected_output = ( + 'GremlinScenarioParallelNode({"name": "a mock parallel node"})' + ) + kwargs = { + "name": "a mock parallel node", + } + helper = GremlinScenarioParallelNode(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + + def test_gremlin_scenario_attack_node_repr_str(self) -> None: + expected_output = 'GremlinScenarioAttackNode({"name": "a mock attack node"})' + kwargs = { + "name": "a mock attack node", + } + helper = GremlinScenarioAttackNode(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) def test_gremlin_scenario_ilfi_node_repr_node(self) -> None: helper = GremlinScenarioILFINode(**mock_ilfi_node) @@ -301,7 +369,64 @@ def test_gremlin_scenario_ilfi_node_repr_node(self) -> None: "type": "InfraAttack", } - self.assertEqual(helper.repr_model(), expected_output) + self.assertEqual(helper.api_model(), expected_output) + + def test_gremlin_scenario_ilfi_node_repr_str(self) -> None: + expected_output = "GremlinScenarioILFINode({'name': 'mock_scenario', 'command': 'GremlinAttackCommandHelper({\"length\": 70})', 'target': 'GremlinAttackTargetHelper({\"exact\": 0, \"percent\": 15, \"strategy_type\": \"Random\"})'})" + kwargs_ch = {"length": 70} + kwargs_th = {"exact": 0, "percent": 15, "strategy_type": "Random"} + kwargs = { + "name": "mock_scenario", + "command": GremlinAttackCommandHelper(**kwargs_ch), + "target": GremlinAttackTargetHelper(**kwargs_th), + } + helper = GremlinScenarioILFINode(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + + # def test_gremlin_scenario_alfi_node_repr_str(self) -> None: + # expected_output = 'GremlinScenarioALFINode({"name": "a mock alfi node"})' + # kwargs = { + # "name": "a mock alfi node", + # } + # helper = GremlinScenarioALFINode(**kwargs) + # helper_output = repr(helper) + # self.assertEqual(expected_output, helper_output) + # helper_output = str(helper) + # self.assertEqual(expected_output, helper_output) + + def test_gremlin_scenario_delay_node_repr_str(self) -> None: + expected_output = ( + 'GremlinScenarioDelayNode({"name": "a mock delay node", "delay": 100})' + ) + kwargs = { + "name": "a mock delay node", + "delay": 100, + } + helper = GremlinScenarioDelayNode(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) + + def test_gremlin_scenario_status_check_node_repr_str(self) -> None: + expected_output = 'GremlinScenarioStatusCheckNode({"name": "a mock status check node", "description": "A mock description", "endpoint_url": "www.example.com", "endpoint_headers": {}, "evaluation_ok_status_codes": ["300-400"], "evaluation_ok_latency_max": 1000, "evaluation_response_body_evaluation": "a mock resposne body"})' + kwargs = { + "name": "a mock status check node", + "description": "A mock description", + "endpoint_url": "www.example.com", + "endpoint_headers": {}, + "evaluation_ok_status_codes": ["300-400"], + "evaluation_ok_latency_max": 1000, + "evaluation_response_body_evaluation": "a mock resposne body", + } + helper = GremlinScenarioStatusCheckNode(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) def test_gremlin_scenario_delay_node_repr_node(self) -> None: helper = GremlinScenarioDelayNode(**mock_delay_node) @@ -313,9 +438,9 @@ def test_gremlin_scenario_delay_node_repr_node(self) -> None: "type": "Delay", } - self.assertEqual(helper.repr_model(), expected_output) + self.assertEqual(helper.api_model(), expected_output) - def test_gremlin_scenario_status_check_node_repr_model(self) -> None: + def test_gremlin_scenario_status_check_node_api_model(self) -> None: helper = GremlinScenarioStatusCheckNode(**mock_status_check_node) expected_output = { "endpointConfiguration": { @@ -335,7 +460,16 @@ def test_gremlin_scenario_status_check_node_repr_model(self) -> None: "thirdPartyPresets": "PythonSDK", "type": "SynchronousStatusCheck", } - self.assertEqual(helper.repr_model(), expected_output) + self.assertEqual(helper.api_model(), expected_output) + + def test_node_graph_repr_str(self) -> None: + expected_output = "_GremlinNodeGraph()" + kwargs = {} + helper = _GremlinNodeGraph(**kwargs) + helper_output = repr(helper) + self.assertEqual(expected_output, helper_output) + helper_output = str(helper) + self.assertEqual(expected_output, helper_output) def test_node_graph_add_edge(self) -> None: helper = _GremlinNodeGraph()