Miguel commited on
Commit
b0da92f
·
1 Parent(s): dd8af64

feat: add cache to mittre retrieval

Browse files
tdagent/tools/retrieve_from_mitre_attack.py CHANGED
@@ -1,9 +1,17 @@
1
  from typing import Any
2
 
 
3
  import gradio as gr
4
  from attackcti import attack_client
5
 
6
 
 
 
 
 
 
 
 
7
  def get_stix_object_of_attack_id(
8
  attack_id: str,
9
  object_type: str = "attack-pattern",
 
1
  from typing import Any
2
 
3
+ import cachetools
4
  import gradio as gr
5
  from attackcti import attack_client
6
 
7
 
8
+ _CACHE_MAX_SIZE = 4096
9
+ _CACHE_TTL_SECONDS = 3600
10
+
11
+
12
+ @cachetools.cached(
13
+ cache=cachetools.TTLCache(maxsize=_CACHE_MAX_SIZE, ttl=_CACHE_TTL_SECONDS),
14
+ )
15
  def get_stix_object_of_attack_id(
16
  attack_id: str,
17
  object_type: str = "attack-pattern",