@cryptotaxi247 / CoPilot
Code Commits Issues Pulls Sessions CI/CD
main
main
py 166 lines 5.54 KB
Raw
1 from typing import Dict
2 from typing import List
3
4 from fastapi import APIRouter
5 from fastapi import Security
6 from loguru import logger
7
8 from app.auth.utils import AuthHandler
9 from app.connectors.graylog.schema.pipelines import GraylogPipelinesResponse
10 from app.connectors.graylog.schema.pipelines import GraylogPipelinesResponseWithRuleID
11 from app.connectors.graylog.schema.pipelines import Pipeline
12 from app.connectors.graylog.schema.pipelines import PipelineRule
13 from app.connectors.graylog.schema.pipelines import PipelineRulesResponse
14 from app.connectors.graylog.schema.pipelines import PipelineWithRuleID
15 from app.connectors.graylog.schema.pipelines import Stage
16 from app.connectors.graylog.schema.pipelines import StageWithRuleID
17 from app.connectors.graylog.services.pipelines import get_pipeline_rule_by_id
18 from app.connectors.graylog.services.pipelines import get_pipeline_rules
19 from app.connectors.graylog.services.pipelines import get_pipelines
20
21 # App specific imports
22
23
24 graylog_pipelines_router = APIRouter()
25
26
27 def create_rule_title_to_id_dict(pipeline_rules: List[PipelineRule]) -> Dict[str, str]:
28 """
29 Creates a dictionary mapping rule titles to rule IDs.
30
31 Args:
32 pipeline_rules (List[PipelineRule]): List of pipeline rules.
33
34 Returns:
35 Dict[str, str]: Dictionary mapping rule titles to rule IDs.
36 """
37 rule_title_to_id = {}
38 for rule in pipeline_rules:
39 rule_title_to_id[rule.title] = rule.id
40 return rule_title_to_id
41
42
43 def transform_stages_with_rule_ids(
44 stages: List[Stage],
45 rule_title_to_id: Dict[str, str],
46 ) -> List[StageWithRuleID]:
47 """
48 Transforms a list of stages by adding corresponding rule IDs based on a dictionary mapping rule titles to IDs.
49
50 Args:
51 stages (List[Stage]): The list of stages to transform.
52 rule_title_to_id (Dict[str, str]): The dictionary mapping rule titles to IDs.
53
54 Returns:
55 List[StageWithRuleID]: The transformed list of stages with added rule IDs.
56 """
57 new_stages = []
58 for stage in stages:
59 rule_ids = [rule_title_to_id.get(rule_title, None) for rule_title in stage.rules]
60 new_stage = StageWithRuleID(**stage.model_dump(), rule_ids=rule_ids)
61 new_stages.append(new_stage)
62 return new_stages
63
64
65 def transform_pipeline_with_rule_ids(
66 pipeline: Pipeline,
67 rule_title_to_id: Dict[str, str],
68 ) -> PipelineWithRuleID:
69 """
70 Transforms a pipeline by replacing rule titles with rule IDs.
71
72 Args:
73 pipeline (Pipeline): The original pipeline object.
74 rule_title_to_id (Dict[str, str]): A dictionary mapping rule titles to rule IDs.
75
76 Returns:
77 PipelineWithRuleID: The transformed pipeline object with rule IDs.
78
79 """
80 new_stages = transform_stages_with_rule_ids(pipeline.stages, rule_title_to_id)
81 pipeline_dict = pipeline.model_dump()
82 pipeline_dict["stages"] = new_stages
83 return PipelineWithRuleID(**pipeline_dict)
84
85
86 @graylog_pipelines_router.get(
87 "/pipelines",
88 response_model=GraylogPipelinesResponse,
89 description="Get all pipelines",
90 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],
91 )
92 async def get_all_pipelines() -> GraylogPipelinesResponse:
93 """
94 Get all pipelines.
95
96 Returns:
97 GraylogPipelinesResponse: The response model containing the pipelines.
98 """
99 logger.info("Fetching all graylog pipelines")
100 return await get_pipelines()
101
102
103 @graylog_pipelines_router.get(
104 "/pipeline/full",
105 response_model=GraylogPipelinesResponseWithRuleID,
106 description="Get all pipelines with rule IDs",
107 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],
108 )
109 async def get_all_pipelines_with_rule_ids() -> GraylogPipelinesResponseWithRuleID:
110 """
111 Retrieve all pipelines with their associated rule IDs.
112
113 Returns:
114 GraylogPipelinesResponseWithRuleID: The response containing the pipelines with rule IDs.
115 """
116 pipelines_response = await get_pipelines()
117 pipeline_rules_response = await get_pipeline_rules()
118
119 rule_title_to_id = create_rule_title_to_id_dict(
120 pipeline_rules_response.pipeline_rules,
121 )
122
123 new_pipelines = [transform_pipeline_with_rule_ids(pipeline, rule_title_to_id) for pipeline in pipelines_response.pipelines]
124
125 return GraylogPipelinesResponseWithRuleID(
126 pipelines=new_pipelines,
127 success=pipelines_response.success,
128 message=pipelines_response.message,
129 )
130
131
132 @graylog_pipelines_router.get(
133 "/pipeline/rules",
134 response_model=PipelineRulesResponse,
135 description="Get all pipeline rules",
136 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],
137 )
138 async def get_all_pipeline_rules() -> PipelineRulesResponse:
139 """
140 Fetches all graylog pipeline rules.
141
142 Returns:
143 PipelineRulesResponse: The response containing all pipeline rules.
144 """
145 logger.info("Fetching all graylog pipeline rules")
146 return await get_pipeline_rules()
147
148
149 @graylog_pipelines_router.get(
150 "/pipeline/rules/{pipeline_id}",
151 response_model=PipelineRulesResponse,
152 description="Get all pipeline rules for a pipeline",
153 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],
154 )
155 async def get_pipeline_rules_for_pipeline(pipeline_id: str) -> PipelineRulesResponse:
156 """
157 Get all pipeline rules for a specific pipeline.
158
159 Args:
160 pipeline_id (str): The ID of the pipeline.
161
162 Returns:
163 PipelineRulesResponse: The response containing the pipeline rules.
164 """
165 logger.info(f"Fetching all graylog pipeline rules for pipeline {pipeline_id}")
166 return await get_pipeline_rule_by_id(pipeline_id)