| 47 | |
| 48 | |
| 49 | class Ability(FirstClassObjectInterface, BaseObject): |
| 50 | |
| 51 | schema = AbilitySchema() |
| 52 | display_schema = AbilitySchema() |
| 53 | |
| 54 | HOOKS = dict() |
| 55 | |
| 56 | @property |
| 57 | def unique(self): |
| 58 | return self.ability_id |
| 59 | |
| 60 | @property |
| 61 | def executors(self): |
| 62 | yield from self._executor_map.values() |
| 63 | |
| 64 | def __init__(self, ability_id='', name=None, description=None, tactic=None, technique_id=None, technique_name=None, |
| 65 | executors=(), requirements=None, privilege=None, repeatable=False, buckets=None, access=None, |
| 66 | additional_info=None, tags=None, singleton=False, plugin='', delete_payload=True, **kwargs): |
| 67 | super().__init__() |
| 68 | self.ability_id = ability_id if ability_id else str(uuid.uuid4()) |
| 69 | self.tactic = tactic.lower() if tactic else None |
| 70 | self.technique_name = technique_name |
| 71 | self.technique_id = technique_id |
| 72 | self.name = name |
| 73 | self.description = description |
| 74 | |
| 75 | self._executor_map = collections.OrderedDict() |
| 76 | self.add_executors(executors) |
| 77 | |
| 78 | self.requirements = requirements if requirements else [] |
| 79 | self.privilege = privilege |
| 80 | self.repeatable = repeatable |
| 81 | self.buckets = buckets if buckets else [] |
| 82 | self.singleton = singleton |
| 83 | if access: |
| 84 | self.access = self.Access(access) |
| 85 | self.additional_info = additional_info or dict() |
| 86 | self.additional_info.update(**kwargs) |
| 87 | self.tags = set(tags) if tags else set() |
| 88 | self.plugin = plugin |
| 89 | self.delete_payload = delete_payload |
| 90 | |
| 91 | def __getattr__(self, item): |
| 92 | try: |
| 93 | return super().__getattribute__('additional_info')[item] |
| 94 | except KeyError: |
| 95 | raise AttributeError(item) |
| 96 | |
| 97 | def store(self, ram): |
| 98 | existing = self.retrieve(ram['abilities'], self.unique) |
| 99 | if not existing: |
| 100 | name_match = [x for x in ram['abilities'] if x.name == self.name] |
| 101 | if name_match: |
| 102 | self.name = self.name + " (2)" |
| 103 | logging.debug(f"Collision in ability name detected for {self.ability_id} and {name_match[0].ability_id} " |
| 104 | f"({name_match[0].name}). Modifying name of the second ability to {self.name}...") |
| 105 | ram['abilities'].append(self) |
| 106 | return self.retrieve(ram['abilities'], self.unique) |