@cryptotaxi247 / CoPilot
Code Commits Issues Pulls Sessions CI/CD
main
main
py 177 lines 5.82 KB
Raw
1 from datetime import datetime
2 from typing import List
3 from typing import Optional
4
5 from pydantic import BaseModel
6 from pydantic import ConfigDict
7 from pydantic import Field
8
9
10 class CustomerRequestBody(BaseModel):
11 customer_code: str = Field(..., description="Unique code for the customer")
12 customer_name: str = Field(..., description="Name of the customer")
13 contact_last_name: str = Field(..., description="Last name of the contact person")
14 contact_first_name: str = Field(..., description="First name of the contact person")
15
16 parent_customer_code: Optional[str] = Field(
17 None,
18 description="Code for the parent customer",
19 )
20 phone: Optional[str] = Field(None, description="Phone number")
21 address_line1: Optional[str] = Field(None, description="First line of the address")
22 address_line2: Optional[str] = Field(None, description="Second line of the address")
23 city: Optional[str] = Field(None, description="City")
24 state: Optional[str] = Field(None, description="State")
25 postal_code: Optional[str] = Field(None, description="Postal Code")
26 country: Optional[str] = Field(None, description="Country")
27 customer_type: Optional[str] = Field(None, description="Type of the customer")
28 logo_file: Optional[str] = Field(None, description="Logo file for the customer")
29 is_provisioned: Optional[bool] = Field(None, description="Whether the customer has been provisioned")
30 model_config = ConfigDict(
31 from_attributes=True,
32 json_schema_extra={
33 "example": {
34 "customer_code": "CUST123",
35 "customer_name": "Sample Customer",
36 "contact_last_name": "Doe",
37 "contact_first_name": "John",
38 "phone": "123-456-7890",
39 "address_line1": "123 Main St",
40 "address_line2": "Apt 4",
41 "city": "Anytown",
42 "state": "CA",
43 "postal_code": "12345",
44 "country": "USA",
45 "customer_type": "Enterprise",
46 "logo_file": "logo.png",
47 "is_provisioned": True,
48 },
49 },
50 )
51
52
53 class CustomerResponse(BaseModel):
54 customer: Optional[CustomerRequestBody] = None
55 success: bool
56 message: str
57
58
59 class CustomersResponse(BaseModel):
60 customers: list[CustomerRequestBody]
61 success: bool
62 message: str
63
64
65 ############# Customer Meta
66 class CustomerMetaRequestBody(BaseModel):
67 customer_meta_graylog_index: str = Field(
68 ...,
69 description="Graylog index for the customer",
70 )
71 customer_meta_graylog_stream: str = Field(
72 ...,
73 description="Graylog stream for the customer",
74 )
75 customer_meta_grafana_org_id: str = Field(
76 ...,
77 description="Grafana organization for the customer",
78 )
79 customer_meta_wazuh_group: str = Field(
80 ...,
81 description="Wazuh group for the customer",
82 )
83 customer_meta_index_retention: str = Field(
84 ...,
85 description="Index retention for the customer",
86 )
87 customer_meta_wazuh_registration_port: str = Field(
88 ...,
89 description="Wazuh registration port for the customer",
90 )
91 customer_meta_wazuh_log_ingestion_port: str = Field(
92 ...,
93 description="Wazuh log ingestion port for the customer",
94 )
95 customer_meta_wazuh_api_port: str = Field(
96 ...,
97 description="Wazuh API port for the customer",
98 )
99 customer_meta_wazuh_auth_password: str = Field(
100 ...,
101 description="Wazuh auth password for the customer",
102 )
103 customer_meta_portainer_stack_id: Optional[int] = Field(
104 None,
105 description="Portainer stack ID for the customer",
106 )
107 model_config = ConfigDict(
108 from_attributes=True,
109 json_schema_extra={
110 "example": {
111 "customer_meta_graylog_index": "graylog_index",
112 "customer_meta_graylog_stream": "graylog_stream",
113 "customer_meta_grafana_org_id": "grafana_org",
114 "customer_meta_wazuh_group": "wazuh_group",
115 "customer_meta_index_retention": "30D",
116 "customer_meta_wazuh_registration_port": "1514",
117 "customer_meta_wazuh_log_ingestion_port": "1515",
118 "customer_meta_wazuh_auth_password": "wazuh_password",
119 },
120 },
121 )
122
123
124 class CustomerMetaResponse(BaseModel):
125 customer_meta: Optional[CustomerMetaRequestBody] = None
126 success: bool
127 message: str
128
129
130 ############# Customer Full Response
131 class CustomerFullResponse(BaseModel):
132 customer: Optional[CustomerRequestBody] = None
133 customer_meta: Optional[CustomerMetaRequestBody] = None
134 success: bool
135 message: str
136
137
138 ############# Agent Model #############
139 class AgentModel(BaseModel):
140 id: Optional[int] = None
141 os: Optional[str] = None
142 label: Optional[str] = None
143 wazuh_last_seen: Optional[datetime] = None
144 velociraptor_last_seen: Optional[datetime] = None
145 velociraptor_agent_version: Optional[str] = None
146 ip_address: Optional[str] = None
147 agent_id: Optional[str] = None
148 hostname: Optional[str] = None
149 critical_asset: Optional[bool] = None
150 velociraptor_id: Optional[str] = None
151 wazuh_agent_version: Optional[str] = None
152 customer_code: Optional[str] = None
153 model_config = ConfigDict(from_attributes=True)
154
155
156 class AgentsResponse(BaseModel):
157 agents: Optional[List[AgentModel]] = Field([], description="List of agents")
158 success: bool
159 message: str
160
161
162 class DeleteCustomerResponse(BaseModel):
163 """
164 Response model for customer deletion.
165 """
166
167 success: bool
168 message: str
169 model_config = ConfigDict(
170 from_attributes=True,
171 json_schema_extra={
172 "example": {
173 "success": True,
174 "message": "Customer 'customer_code' deleted successfully",
175 },
176 },
177 )