Convert to mysql (#195)
* convert to mysql initial push * Update database connection and create copilot user * Update MySQL database URIs in db_session.py * Update dependency in docker-compose.yml * added wait-fit-it script during startup * logging for debugging * Update database configuration and module description * Add environment variable for server IP * fix alembic versioning * Add ADMIN_USER_NEEDED environment variable and conditionally create admin user * Remove unused ADMIN_USER_NEEDED variable and update database initialization * precommit fixes for conversion to mysql * Update Docker tags for backend and frontend images
taylor_socfortress committed
Apr 23, 2024 at 09:08 UTC
f22d0b4ac9080003f68d2c92e9d759ccc4882e86
27 files changed
+1533
-223
.env.example
+5
@@ -1,6 +1,11 @@
1
# Leave this as is if connecting from a remote machine
2
SERVER_IP=0.0.0.0
3
4
+MYSQL_URL=copilot-mysql
5
+MYSQL_ROOT_PASSWORD=REPLACE_WITH_PASSWORD
6
+MYSQL_USER=copilot
7
+MYSQL_PASSWORD=REPLACE_WITH_PASSWORD
8
+
9
# ! ALERT FORWARDING IP
10
# Set this to the IP of the host running CoPilot. This is used by Graylog to forward alerts to CoPilot
11
# ! Ensure Graylog is able to reach this IP and port 5000
.flake8
+2
-1
@@ -3,7 +3,7 @@
3
max-line-length = 180
4
#select = B,C,E,F,W,T4,B9
5
#ignore = E203, E266, E501, W503, F403, F401
6
-ignore = E402, W503, E231, W605, E266
6
+ignore = E402, W503, E231, W605, E266, E712
7
# E402, # module level import not at top of file (using isort)
8
# W503, # line break before binary operator
9
# E231, # missing whitespace after ',' (caused by black style)
@@ -11,3 +11,4 @@ ignore = E402, W503, E231, W605, E266
11
extend-ignore = E203
12
exclude =
13
.venv
14
+ backend/alembic/env.py
backend/Dockerfile
+5
-5
@@ -1,7 +1,3 @@
1
-# build with `docker build -t python-backend -f Dockerfile.deb .`
2
-# run with `docker run -p 5000:5000 -d python-backend`
3
-# Start with the base Debian 11 image
4
-# looking to split into 2 containers, one for the backend and one for the frontend
1
FROM debian:11
2
3
# Set environment variables
@@ -51,6 +47,8 @@ RUN apt-get install -y wkhtmltopdf
47
# Copy your application into the Docker image
48
WORKDIR /opt/copilot/backend
49
COPY . .
50
+COPY wait-for-it.sh /usr/wait-for-it.sh
51
+RUN chmod +x /usr/wait-for-it.sh
52
# Create file-store folder
53
RUN mkdir file-store
54
@@ -113,4 +111,6 @@ ARG COPILOT_API_KEY
111
ENV COPILOT_API_KEY=$COPILOT_API_KEY
112
113
# Run your application
116
-CMD ["sh", "-c", "ls -la && /opt/venv/bin/python copilot.py"]
114
+# CMD ["sh", "-c", "ls -la && /opt/venv/bin/python copilot.py"]
115
+# Use wait-for-it.sh to wait for the MySQL service to be ready before starting your application
116
+CMD ["/usr/wait-for-it.sh", "copilot-mysql:3306", "--", "/opt/venv/bin/python", "copilot.py"]
backend/alembic/README
new
+1
@@ -0,0 +1 @@
1
+Generic single-database configuration.
backend/alembic/alembic.ini
new
+116
@@ -0,0 +1,116 @@
1
+# A generic, single database configuration.
2
+
3
+[alembic]
4
+# path to migration scripts
5
+script_location = alembic
6
+
7
+# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s
8
+# Uncomment the line below if you want the files to be prepended with date and time
9
+# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file
10
+# for all available tokens
11
+# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s
12
+
13
+# sys.path path, will be prepended to sys.path if present.
14
+# defaults to the current working directory.
15
+prepend_sys_path = .
16
+
17
+# timezone to use when rendering the date within the migration file
18
+# as well as the filename.
19
+# If specified, requires the python>=3.9 or backports.zoneinfo library.
20
+# Any required deps can installed by adding `alembic[tz]` to the pip requirements
21
+# string value is passed to ZoneInfo()
22
+# leave blank for localtime
23
+# timezone =
24
+
25
+# max length of characters to apply to the
26
+# "slug" field
27
+# truncate_slug_length = 40
28
+
29
+# set to 'true' to run the environment during
30
+# the 'revision' command, regardless of autogenerate
31
+# revision_environment = false
32
+
33
+# set to 'true' to allow .pyc and .pyo files without
34
+# a source .py file to be detected as revisions in the
35
+# versions/ directory
36
+# sourceless = false
37
+
38
+# version location specification; This defaults
39
+# to alembic/versions. When using multiple version
40
+# directories, initial revisions must be specified with --version-path.
41
+# The path separator used here should be the separator specified by "version_path_separator" below.
42
+# version_locations = %(here)s/bar:%(here)s/bat:alembic/versions
43
+
44
+# version path separator; As mentioned above, this is the character used to split
45
+# version_locations. The default within new alembic.ini files is "os", which uses os.pathsep.
46
+# If this key is omitted entirely, it falls back to the legacy behavior of splitting on spaces and/or commas.
47
+# Valid values for version_path_separator are:
48
+#
49
+# version_path_separator = :
50
+# version_path_separator = ;
51
+# version_path_separator = space
52
+version_path_separator = os # Use os.pathsep. Default configuration used for new projects.
53
+
54
+# set to 'true' to search source files recursively
55
+# in each "version_locations" directory
56
+# new in Alembic version 1.10
57
+# recursive_version_locations = false
58
+
59
+# the output encoding used when revision files
60
+# are written from script.py.mako
61
+# output_encoding = utf-8
62
+
63
+sqlalchemy.url = mysql+pymysql://copilot:REPLACE_WITH_PASS@copilot-mysql/copilot
64
+
65
+
66
+[post_write_hooks]
67
+# post_write_hooks defines scripts or Python functions that are run
68
+# on newly generated revision scripts. See the documentation for further
69
+# detail and examples
70
+
71
+# format using "black" - use the console_scripts runner, against the "black" entrypoint
72
+# hooks = black
73
+# black.type = console_scripts
74
+# black.entrypoint = black
75
+# black.options = -l 79 REVISION_SCRIPT_FILENAME
76
+
77
+# lint with attempts to fix using "ruff" - use the exec runner, execute a binary
78
+# hooks = ruff
79
+# ruff.type = exec
80
+# ruff.executable = %(here)s/.venv/bin/ruff
81
+# ruff.options = --fix REVISION_SCRIPT_FILENAME
82
+
83
+# Logging configuration
84
+[loggers]
85
+keys = root,sqlalchemy,alembic
86
+
87
+[handlers]
88
+keys = console
89
+
90
+[formatters]
91
+keys = generic
92
+
93
+[logger_root]
94
+level = WARN
95
+handlers = console
96
+qualname =
97
+
98
+[logger_sqlalchemy]
99
+level = WARN
100
+handlers =
101
+qualname = sqlalchemy.engine
102
+
103
+[logger_alembic]
104
+level = INFO
105
+handlers =
106
+qualname = alembic
107
+
108
+[handler_console]
109
+class = StreamHandler
110
+args = (sys.stderr,)
111
+level = NOTSET
112
+formatter = generic
113
+
114
+[formatter_generic]
115
+format = %(levelname)-5.5s [%(name)s] %(message)s
116
+datefmt = %H:%M:%S
backend/alembic/env.py
new
+99
@@ -0,0 +1,99 @@
1
+from logging.config import fileConfig
2
+
3
+from sqlalchemy import engine_from_config
4
+from sqlalchemy import pool
5
+from sqlmodel import SQLModel
6
+
7
+from alembic import context
8
+
9
+# from app.db.all_models import *
10
+from app.auth.models.users import User
11
+from app.connectors.models import Connectors
12
+
13
+# from app.integrations.sap_siem.models.sap_siem import SapSiemMultipleLogins
14
+from app.customer_provisioning.models.default_settings import (
15
+ CustomerProvisioningDefaultSettings,
16
+)
17
+
18
+# from app.connectors.sublime.models.alerts import SublimeAlerts
19
+# from app.connectors.wazuh_manager.models.rules import DisabledRule
20
+from app.db.universal_models import Agents
21
+from app.db.universal_models import Customers
22
+from app.db.universal_models import CustomersMeta
23
+from app.db.universal_models import LogEntry
24
+from app.integrations.alert_creation_settings.models.alert_creation_settings import (
25
+ AlertCreationSettings,
26
+)
27
+from app.integrations.models.customer_integration_settings import CustomerIntegrations
28
+from app.integrations.monitoring_alert.models.monitoring_alert import MonitoringAlerts
29
+from app.schedulers.models.scheduler import JobMetadata
30
+
31
+# this is the Alembic Config object, which provides
32
+# access to the values within the .ini file in use.
33
+config = context.config
34
+
35
+# Interpret the config file for Python logging.
36
+# This line sets up loggers basically.
37
+if config.config_file_name is not None:
38
+ fileConfig(config.config_file_name)
39
+
40
+# add your model's MetaData object here
41
+# for 'autogenerate' support
42
+# from myapp import mymodel
43
+# target_metadata = mymodel.Base.metadata
44
+target_metadata = SQLModel.metadata
45
+
46
+# other values from the config, defined by the needs of env.py,
47
+# can be acquired:
48
+# my_important_option = config.get_main_option("my_important_option")
49
+# ... etc.
50
+
51
+
52
+def run_migrations_offline() -> None:
53
+ """Run migrations in 'offline' mode.
54
+
55
+ This configures the context with just a URL
56
+ and not an Engine, though an Engine is acceptable
57
+ here as well. By skipping the Engine creation
58
+ we don't even need a DBAPI to be available.
59
+
60
+ Calls to context.execute() here emit the given string to the
61
+ script output.
62
+
63
+ """
64
+ url = config.get_main_option("sqlalchemy.url")
65
+ context.configure(
66
+ url=url,
67
+ target_metadata=target_metadata,
68
+ literal_binds=True,
69
+ dialect_opts={"paramstyle": "named"},
70
+ )
71
+
72
+ with context.begin_transaction():
73
+ context.run_migrations()
74
+
75
+
76
+def run_migrations_online() -> None:
77
+ """Run migrations in 'online' mode.
78
+
79
+ In this scenario we need to create an Engine
80
+ and associate a connection with the context.
81
+
82
+ """
83
+ connectable = engine_from_config(
84
+ config.get_section(config.config_ini_section, {}),
85
+ prefix="sqlalchemy.",
86
+ poolclass=pool.NullPool,
87
+ )
88
+
89
+ with connectable.connect() as connection:
90
+ context.configure(connection=connection, target_metadata=target_metadata)
91
+
92
+ with context.begin_transaction():
93
+ context.run_migrations()
94
+
95
+
96
+if context.is_offline_mode():
97
+ run_migrations_offline()
98
+else:
99
+ run_migrations_online()
backend/alembic/script.py.mako
new
+26
@@ -0,0 +1,26 @@
1
+"""${message}
2
+
3
+Revision ID: ${up_revision}
4
+Revises: ${down_revision | comma,n}
5
+Create Date: ${create_date}
6
+
7
+"""
8
+from typing import Sequence, Union
9
+
10
+from alembic import op
11
+import sqlalchemy as sa
12
+${imports if imports else ""}
13
+
14
+# revision identifiers, used by Alembic.
15
+revision: str = ${repr(up_revision)}
16
+down_revision: Union[str, None] = ${repr(down_revision)}
17
+branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
18
+depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
19
+
20
+
21
+def upgrade() -> None:
22
+ ${upgrades if upgrades else "pass"}
23
+
24
+
25
+def downgrade() -> None:
26
+ ${downgrades if downgrades else "pass"}
backend/alembic/versions/bdf40d064ed1_initial_database_migration.py
new
+403
@@ -0,0 +1,403 @@
1
+"""Initial database migration
2
+
3
+Revision ID: bdf40d064ed1
4
+Revises:
5
+Create Date: 2024-04-22 19:48:15.129902
6
+
7
+"""
8
+from typing import Sequence
9
+from typing import Union
10
+
11
+import sqlalchemy as sa
12
+import sqlalchemy.dialects.mysql as mysql
13
+
14
+from alembic import op
15
+
16
+# revision identifiers, used by Alembic.
17
+revision: str = "bdf40d064ed1"
18
+down_revision: Union[str, None] = None
19
+branch_labels: Union[str, Sequence[str], None] = None
20
+depends_on: Union[str, Sequence[str], None] = None
21
+
22
+
23
+def upgrade() -> None:
24
+ # ### commands auto generated by Alembic - please adjust! ###
25
+ op.create_table(
26
+ "available_integrations",
27
+ sa.Column("id", sa.Integer(), nullable=False),
28
+ sa.Column("integration_name", sa.String(length=255), nullable=False),
29
+ sa.Column("description", sa.String(length=1024), nullable=False),
30
+ sa.Column("integration_details", mysql.TEXT(length=1000000), nullable=False),
31
+ sa.PrimaryKeyConstraint("id"),
32
+ )
33
+ op.create_table(
34
+ "connectors",
35
+ sa.Column("id", sa.Integer(), nullable=False),
36
+ sa.Column("connector_name", sa.String(length=256), nullable=False),
37
+ sa.Column("connector_type", sa.String(length=256), nullable=False),
38
+ sa.Column("connector_url", sa.String(length=750), nullable=False),
39
+ sa.Column("connector_last_updated", sa.DateTime(), nullable=False),
40
+ sa.Column("connector_username", sa.String(length=256), nullable=True),
41
+ sa.Column("connector_password", sa.String(length=256), nullable=True),
42
+ sa.Column("connector_api_key", sa.String(length=750), nullable=True),
43
+ sa.Column("connector_description", sa.String(length=1000), nullable=True),
44
+ sa.Column("connector_supports", sa.String(length=1000), nullable=True),
45
+ sa.Column("connector_configured", sa.Boolean(), nullable=False),
46
+ sa.Column("connector_verified", sa.Boolean(), nullable=False),
47
+ sa.Column("connector_accepts_host_only", sa.Boolean(), nullable=False),
48
+ sa.Column("connector_accepts_api_key", sa.Boolean(), nullable=False),
49
+ sa.Column("connector_accepts_username_password", sa.Boolean(), nullable=False),
50
+ sa.Column("connector_accepts_file", sa.Boolean(), nullable=False),
51
+ sa.Column("connector_accepts_extra_data", sa.Boolean(), nullable=False),
52
+ sa.Column("connector_extra_data", sa.String(length=1000), nullable=True),
53
+ sa.Column("connector_enabled", sa.Boolean(), nullable=False),
54
+ sa.PrimaryKeyConstraint("id"),
55
+ )
56
+ op.create_table(
57
+ "custom_alert_creation_settings",
58
+ sa.Column("id", sa.Integer(), nullable=False),
59
+ sa.Column("customer_code", sa.String(length=50), nullable=False),
60
+ sa.Column("customer_name", sa.String(length=50), nullable=False),
61
+ sa.Column("excluded_wazuh_rules", sa.String(length=1024), nullable=True),
62
+ sa.Column("excluded_suricata_rules", sa.String(length=1024), nullable=True),
63
+ sa.Column("timefield", sa.String(length=1024), nullable=True),
64
+ sa.Column("office365_organization_id", sa.String(length=1024), nullable=True),
65
+ sa.Column("iris_customer_id", sa.Integer(), nullable=True),
66
+ sa.Column("iris_customer_name", sa.String(length=1024), nullable=True),
67
+ sa.Column("iris_index", sa.String(length=1024), nullable=True),
68
+ sa.Column("grafana_url", sa.String(length=1024), nullable=True),
69
+ sa.Column("misp_url", sa.String(length=1024), nullable=True),
70
+ sa.Column("opencti_url", sa.String(length=1024), nullable=True),
71
+ sa.Column("custom_message", sa.String(length=1024), nullable=True),
72
+ sa.Column("shuffle_endpoint", sa.String(length=1024), nullable=True),
73
+ sa.Column("nvd_url", sa.String(length=1024), nullable=True),
74
+ sa.PrimaryKeyConstraint("id"),
75
+ )
76
+ op.create_table(
77
+ "customer_integrations",
78
+ sa.Column("id", sa.Integer(), nullable=False),
79
+ sa.Column("customer_code", sa.String(length=50), nullable=False),
80
+ sa.Column("customer_name", sa.String(length=255), nullable=False),
81
+ sa.Column("integration_service_id", sa.Integer(), nullable=False),
82
+ sa.Column("integration_service_name", sa.String(length=255), nullable=False),
83
+ sa.Column("deployed", sa.Boolean(), nullable=False),
84
+ sa.PrimaryKeyConstraint("id"),
85
+ )
86
+ op.create_table(
87
+ "customer_integrations_meta",
88
+ sa.Column("id", sa.Integer(), nullable=False),
89
+ sa.Column("customer_code", sa.String(length=50), nullable=False),
90
+ sa.Column("integration_name", sa.String(length=255), nullable=False),
91
+ sa.Column("graylog_input_id", sa.String(length=1024), nullable=True),
92
+ sa.Column("graylog_index_id", sa.String(length=1024), nullable=False),
93
+ sa.Column("graylog_stream_id", sa.String(length=1024), nullable=False),
94
+ sa.Column("grafana_org_id", sa.String(length=1024), nullable=False),
95
+ sa.Column("grafana_dashboard_folder_id", sa.String(length=1024), nullable=False),
96
+ sa.PrimaryKeyConstraint("id"),
97
+ )
98
+ op.create_table(
99
+ "customer_provisioning_default_settings",
100
+ sa.Column("id", sa.Integer(), nullable=False),
101
+ sa.Column("cluster_name", sa.String(length=50), nullable=False),
102
+ sa.Column("cluster_key", sa.String(length=1000), nullable=False),
103
+ sa.Column("master_ip", sa.String(length=50), nullable=False),
104
+ sa.Column("grafana_url", sa.String(length=1024), nullable=False),
105
+ sa.Column("wazuh_worker_hostname", sa.String(length=100), nullable=False),
106
+ sa.PrimaryKeyConstraint("id"),
107
+ )
108
+ op.create_table(
109
+ "customers",
110
+ sa.Column("id", sa.Integer(), nullable=False),
111
+ sa.Column("customer_code", sa.String(length=50), nullable=False),
112
+ sa.Column("parent_customer_code", sa.String(length=11), nullable=True),
113
+ sa.Column("customer_name", sa.String(length=50), nullable=False),
114
+ sa.Column("contact_last_name", sa.String(length=50), nullable=True),
115
+ sa.Column("contact_first_name", sa.String(length=50), nullable=True),
116
+ sa.Column("phone", sa.String(length=50), nullable=True),
117
+ sa.Column("address_line1", sa.String(length=1024), nullable=True),
118
+ sa.Column("address_line2", sa.String(length=1024), nullable=True),
119
+ sa.Column("city", sa.String(length=50), nullable=True),
120
+ sa.Column("state", sa.String(length=50), nullable=True),
121
+ sa.Column("postal_code", sa.String(length=15), nullable=True),
122
+ sa.Column("country", sa.String(length=50), nullable=True),
123
+ sa.Column("customer_type", sa.String(length=50), nullable=True),
124
+ sa.Column("logo_file", sa.String(length=64), nullable=True),
125
+ sa.Column("created_at", sa.DateTime(), nullable=False),
126
+ sa.PrimaryKeyConstraint("id"),
127
+ )
128
+ op.create_index(op.f("ix_customers_customer_code"), "customers", ["customer_code"], unique=False)
129
+ op.create_table(
130
+ "integration_services",
131
+ sa.Column("id", sa.Integer(), nullable=False),
132
+ sa.Column("service_name", sa.String(length=255), nullable=False),
133
+ sa.Column("auth_type", sa.String(length=50), nullable=False),
134
+ sa.PrimaryKeyConstraint("id"),
135
+ )
136
+ op.create_table(
137
+ "license",
138
+ sa.Column("id", sa.Integer(), nullable=False),
139
+ sa.Column("license_key", sa.String(length=1024), nullable=False),
140
+ sa.Column("customer_name", sa.String(length=1024), nullable=False),
141
+ sa.Column("customer_email", sa.String(length=1024), nullable=False),
142
+ sa.Column("company_name", sa.String(length=1024), nullable=False),
143
+ sa.PrimaryKeyConstraint("id"),
144
+ )
145
+ op.create_table(
146
+ "log_entries",
147
+ sa.Column("id", sa.Integer(), nullable=False),
148
+ sa.Column("timestamp", sa.DateTime(), nullable=False),
149
+ sa.Column("event_type", sa.String(length=256), nullable=False),
150
+ sa.Column("user_id", sa.Integer(), nullable=True),
151
+ sa.Column("route", sa.String(length=256), nullable=True),
152
+ sa.Column("method", sa.String(length=256), nullable=True),
153
+ sa.Column("status_code", sa.Integer(), nullable=False),
154
+ sa.Column("message", sa.String(length=5024), nullable=True),
155
+ sa.Column("additional_info", sa.String(length=5024), nullable=True),
156
+ sa.PrimaryKeyConstraint("id"),
157
+ )
158
+ op.create_table(
159
+ "monitoring_alerts",
160
+ sa.Column("id", sa.Integer(), nullable=False),
161
+ sa.Column("alert_id", sa.String(length=1024), nullable=False),
162
+ sa.Column("alert_index", sa.String(length=1024), nullable=False),
163
+ sa.Column("customer_code", sa.String(length=50), nullable=False),
164
+ sa.Column("alert_source", sa.String(length=1024), nullable=False),
165
+ sa.PrimaryKeyConstraint("id"),
166
+ )
167
+ op.create_table(
168
+ "role",
169
+ sa.Column("id", sa.Integer(), nullable=False),
170
+ sa.Column("name", sa.String(length=256), nullable=False),
171
+ sa.Column("description", sa.String(length=256), nullable=False),
172
+ sa.PrimaryKeyConstraint("id"),
173
+ )
174
+ op.create_table(
175
+ "scheduled_job_metadata",
176
+ sa.Column("id", sa.Integer(), nullable=False),
177
+ sa.Column("job_id", sa.String(length=256), nullable=False),
178
+ sa.Column("last_success", sa.DateTime(), nullable=True),
179
+ sa.Column("time_interval", sa.Integer(), nullable=False),
180
+ sa.Column("extra_data", sa.String(length=256), nullable=True),
181
+ sa.Column("enabled", sa.Boolean(), nullable=False),
182
+ sa.PrimaryKeyConstraint("id"),
183
+ )
184
+ op.create_index(op.f("ix_scheduled_job_metadata_job_id"), "scheduled_job_metadata", ["job_id"], unique=False)
185
+ op.create_table(
186
+ "schedulerjob",
187
+ sa.Column("next_run_time", sa.Float(), nullable=True),
188
+ sa.Column("job_state", sa.LargeBinary(), nullable=False),
189
+ sa.Column("id", sa.String(length=255), nullable=False),
190
+ sa.PrimaryKeyConstraint("id"),
191
+ )
192
+ op.create_index(op.f("ix_schedulerjob_next_run_time"), "schedulerjob", ["next_run_time"], unique=False)
193
+ op.create_table(
194
+ "agents",
195
+ sa.Column("id", sa.Integer(), nullable=False),
196
+ sa.Column("agent_id", sa.String(length=256), nullable=False),
197
+ sa.Column("ip_address", sa.String(length=256), nullable=False),
198
+ sa.Column("os", sa.String(length=256), nullable=False),
199
+ sa.Column("hostname", sa.String(length=256), nullable=False),
200
+ sa.Column("label", sa.String(length=256), nullable=False),
201
+ sa.Column("critical_asset", sa.Boolean(), nullable=False),
202
+ sa.Column("wazuh_last_seen", sa.DateTime(), nullable=False),
203
+ sa.Column("velociraptor_id", sa.String(length=256), nullable=False),
204
+ sa.Column("velociraptor_last_seen", sa.DateTime(), nullable=False),
205
+ sa.Column("wazuh_agent_version", sa.String(length=256), nullable=False),
206
+ sa.Column("wazuh_agent_status", sa.String(length=256), nullable=False),
207
+ sa.Column("velociraptor_agent_version", sa.String(length=256), nullable=False),
208
+ sa.Column("customer_code", sa.String(length=256), nullable=True),
209
+ sa.Column("quarantined", sa.Boolean(), nullable=False),
210
+ sa.ForeignKeyConstraint(
211
+ ["customer_code"],
212
+ ["customers.customer_code"],
213
+ ),
214
+ sa.PrimaryKeyConstraint("id"),
215
+ )
216
+ op.create_index(op.f("ix_agents_agent_id"), "agents", ["agent_id"], unique=False)
217
+ op.create_table(
218
+ "available_integrations_auth_keys",
219
+ sa.Column("id", sa.Integer(), nullable=False),
220
+ sa.Column("integration_id", sa.Integer(), nullable=True),
221
+ sa.Column("integration_name", sa.String(length=255), nullable=False),
222
+ sa.Column("auth_key_name", sa.String(length=255), nullable=False),
223
+ sa.ForeignKeyConstraint(
224
+ ["integration_id"],
225
+ ["available_integrations.id"],
226
+ ),
227
+ sa.PrimaryKeyConstraint("id"),
228
+ )
229
+ op.create_table(
230
+ "connectorhistory",
231
+ sa.Column("id", sa.Integer(), nullable=False),
232
+ sa.Column("connector_id", sa.Integer(), nullable=False),
233
+ sa.Column("change_timestamp", sa.DateTime(), nullable=False),
234
+ sa.Column("change_description", sa.String(length=10000), nullable=False),
235
+ sa.ForeignKeyConstraint(
236
+ ["connector_id"],
237
+ ["connectors.id"],
238
+ ),
239
+ sa.PrimaryKeyConstraint("id"),
240
+ )
241
+ op.create_table(
242
+ "custom_alert_creation_event_order",
243
+ sa.Column("id", sa.Integer(), nullable=False),
244
+ sa.Column("alert_creation_settings_id", sa.Integer(), nullable=True),
245
+ sa.Column("order_label", sa.String(length=255), nullable=False),
246
+ sa.ForeignKeyConstraint(
247
+ ["alert_creation_settings_id"],
248
+ ["custom_alert_creation_settings.id"],
249
+ ),
250
+ sa.PrimaryKeyConstraint("id"),
251
+ )
252
+ op.create_table(
253
+ "customersmeta",
254
+ sa.Column("id", sa.Integer(), nullable=False),
255
+ sa.Column("customer_code", sa.String(length=255), nullable=False),
256
+ sa.Column("customer_name", sa.String(length=255), nullable=False),
257
+ sa.Column("customer_meta_graylog_index", sa.String(length=1024), nullable=False),
258
+ sa.Column("customer_meta_graylog_stream", sa.String(length=1024), nullable=False),
259
+ sa.Column("customer_meta_grafana_org_id", sa.String(length=1024), nullable=False),
260
+ sa.Column("customer_meta_wazuh_group", sa.String(length=1024), nullable=False),
261
+ sa.Column("customer_meta_index_retention", sa.String(length=1024), nullable=True),
262
+ sa.Column("customer_meta_wazuh_registration_port", sa.String(length=1024), nullable=True),
263
+ sa.Column("customer_meta_wazuh_log_ingestion_port", sa.String(length=1024), nullable=True),
264
+ sa.Column("customer_meta_wazuh_api_port", sa.String(length=1024), nullable=True),
265
+ sa.Column("customer_meta_wazuh_auth_password", sa.String(length=1024), nullable=True),
266
+ sa.Column("customer_meta_iris_customer_id", sa.Integer(), nullable=True),
267
+ sa.Column("customer_meta_office365_organization_id", sa.String(length=1024), nullable=True),
268
+ sa.ForeignKeyConstraint(
269
+ ["customer_code"],
270
+ ["customers.customer_code"],
271
+ ),
272
+ sa.PrimaryKeyConstraint("id"),
273
+ )
274
+ op.create_table(
275
+ "integration_configs",
276
+ sa.Column("id", sa.Integer(), nullable=False),
277
+ sa.Column("integration_service_id", sa.Integer(), nullable=True),
278
+ sa.Column("config_key", sa.String(length=255), nullable=False),
279
+ sa.Column("config_value", sa.String(length=1024), nullable=False),
280
+ sa.ForeignKeyConstraint(
281
+ ["integration_service_id"],
282
+ ["integration_services.id"],
283
+ ),
284
+ sa.PrimaryKeyConstraint("id"),
285
+ )
286
+ op.create_table(
287
+ "integration_subscriptions",
288
+ sa.Column("id", sa.Integer(), nullable=False),
289
+ sa.Column("customer_id", sa.Integer(), nullable=True),
290
+ sa.Column("integration_service_id", sa.Integer(), nullable=True),
291
+ sa.ForeignKeyConstraint(
292
+ ["customer_id"],
293
+ ["customer_integrations.id"],
294
+ ),
295
+ sa.ForeignKeyConstraint(
296
+ ["integration_service_id"],
297
+ ["integration_services.id"],
298
+ ),
299
+ sa.PrimaryKeyConstraint("id"),
300
+ )
301
+ op.create_table(
302
+ "user",
303
+ sa.Column("id", sa.Integer(), nullable=False),
304
+ sa.Column("username", sa.String(length=256), nullable=False),
305
+ sa.Column("password", sa.String(length=256), nullable=False),
306
+ sa.Column("email", sa.String(length=1024), nullable=False),
307
+ sa.Column("created_at", sa.DateTime(), nullable=False),
308
+ sa.Column("role_id", sa.Integer(), nullable=True),
309
+ sa.ForeignKeyConstraint(
310
+ ["role_id"],
311
+ ["role.id"],
312
+ ),
313
+ sa.PrimaryKeyConstraint("id"),
314
+ )
315
+ op.create_index(op.f("ix_user_username"), "user", ["username"], unique=False)
316
+ op.create_table(
317
+ "custom_alert_creation_condition",
318
+ sa.Column("id", sa.Integer(), nullable=False),
319
+ sa.Column("event_order_id", sa.Integer(), nullable=True),
320
+ sa.Column("field_name", sa.String(length=1024), nullable=False),
321
+ sa.Column("field_value", sa.String(length=1024), nullable=False),
322
+ sa.ForeignKeyConstraint(
323
+ ["event_order_id"],
324
+ ["custom_alert_creation_event_order.id"],
325
+ ),
326
+ sa.PrimaryKeyConstraint("id"),
327
+ )
328
+ op.create_table(
329
+ "custom_alert_creation_event_config",
330
+ sa.Column("id", sa.Integer(), nullable=False),
331
+ sa.Column("event_order_id", sa.Integer(), nullable=True),
332
+ sa.Column("event_id", sa.String(length=255), nullable=False),
333
+ sa.Column("field", sa.String(length=1024), nullable=False),
334
+ sa.Column("value", sa.String(length=1024), nullable=False),
335
+ sa.ForeignKeyConstraint(
336
+ ["event_order_id"],
337
+ ["custom_alert_creation_event_order.id"],
338
+ ),
339
+ sa.PrimaryKeyConstraint("id"),
340
+ )
341
+ op.create_table(
342
+ "integration_auth_keys",
343
+ sa.Column("id", sa.Integer(), nullable=False),
344
+ sa.Column("subscription_id", sa.Integer(), nullable=True),
345
+ sa.Column("auth_key_name", sa.String(length=255), nullable=False),
346
+ sa.Column("auth_value", sa.String(length=1024), nullable=False),
347
+ sa.ForeignKeyConstraint(
348
+ ["subscription_id"],
349
+ ["integration_subscriptions.id"],
350
+ ),
351
+ sa.PrimaryKeyConstraint("id"),
352
+ )
353
+ op.create_table(
354
+ "smtp",
355
+ sa.Column("id", sa.Integer(), nullable=False),
356
+ sa.Column("email", sa.String(length=1024), nullable=False),
357
+ sa.Column("smtp_password", sa.String(length=256), nullable=False),
358
+ sa.Column("smtp_server", sa.String(length=256), nullable=False),
359
+ sa.Column("smtp_port", sa.Integer(), nullable=False),
360
+ sa.Column("user_id", sa.Integer(), nullable=False),
361
+ sa.ForeignKeyConstraint(
362
+ ["user_id"],
363
+ ["user.id"],
364
+ ),
365
+ sa.PrimaryKeyConstraint("id"),
366
+ )
367
+ # ### end Alembic commands ###
368
+
369
+
370
+def downgrade() -> None:
371
+ # ### commands auto generated by Alembic - please adjust! ###
372
+ op.drop_table("smtp")
373
+ op.drop_table("integration_auth_keys")
374
+ op.drop_table("custom_alert_creation_event_config")
375
+ op.drop_table("custom_alert_creation_condition")
376
+ op.drop_index(op.f("ix_user_username"), table_name="user")
377
+ op.drop_table("user")
378
+ op.drop_table("integration_subscriptions")
379
+ op.drop_table("integration_configs")
380
+ op.drop_table("customersmeta")
381
+ op.drop_table("custom_alert_creation_event_order")
382
+ op.drop_table("connectorhistory")
383
+ op.drop_table("available_integrations_auth_keys")
384
+ op.drop_index(op.f("ix_agents_agent_id"), table_name="agents")
385
+ op.drop_table("agents")
386
+ op.drop_index(op.f("ix_schedulerjob_next_run_time"), table_name="schedulerjob")
387
+ op.drop_table("schedulerjob")
388
+ op.drop_index(op.f("ix_scheduled_job_metadata_job_id"), table_name="scheduled_job_metadata")
389
+ op.drop_table("scheduled_job_metadata")
390
+ op.drop_table("role")
391
+ op.drop_table("monitoring_alerts")
392
+ op.drop_table("log_entries")
393
+ op.drop_table("license")
394
+ op.drop_table("integration_services")
395
+ op.drop_index(op.f("ix_customers_customer_code"), table_name="customers")
396
+ op.drop_table("customers")
397
+ op.drop_table("customer_provisioning_default_settings")
398
+ op.drop_table("customer_integrations_meta")
399
+ op.drop_table("customer_integrations")
400
+ op.drop_table("custom_alert_creation_settings")
401
+ op.drop_table("connectors")
402
+ op.drop_table("available_integrations")
403
+ # ### end Alembic commands ###
backend/app/agents/services/sync.py
+8
-1
@@ -1,5 +1,6 @@
1
from typing import List
2
3
+from fastapi import HTTPException
4
from loguru import logger
5
from sqlalchemy.ext.asyncio import AsyncSession
6
from sqlalchemy.future import select
@@ -67,7 +68,13 @@ async def add_agent_to_db(
68
"""
69
new_agent = Agents.create_from_model(agent, client, customer_code)
70
session.add(new_agent)
70
- await session.commit() # Use the await keyword to commit asynchronously
71
+ logger.info(f"Adding agent {agent.agent_name} to the database")
72
+ try:
73
+ await session.commit() # Use the await keyword to commit asynchronously
74
+ except Exception as e:
75
+ logger.error(f"Failed to add agent {agent.agent_name} to the database: {e}")
76
+ await session.rollback()
77
+ raise HTTPException(status_code=500, detail=str(e))
78
logger.info(f"Agent {agent.agent_name} added to the database")
79
80
backend/app/auth/models/users.py
+7
-1
@@ -24,7 +24,7 @@ class Role(SQLModel, table=True):
24
25
class User(SQLModel, table=True):
26
id: Optional[int] = Field(primary_key=True)
27
- username: str = Field(index=True)
27
+ username: str = Field(index=True, max_length=256)
28
password: str = Field(max_length=256, min_length=6)
29
email: EmailStr
30
created_at: datetime.datetime = datetime.datetime.now()
@@ -56,6 +56,12 @@ class UserInput(SQLModel):
56
foreign_key="role.id",
57
)
58
59
+ @validator("role_id")
60
+ def check_role_id(cls, value):
61
+ if value not in [e.value for e in RoleEnum]:
62
+ raise ValueError("Invalid role ID")
63
+ return value
64
+
65
66
class UserLogin(SQLModel):
67
username: str
backend/app/auth/routes/auth.py
+1
-1
@@ -99,7 +99,7 @@ async def register(user: UserInput, session: AsyncSession = Depends(get_db)):
99
username=user.username,
100
password=hashed_pwd,
101
email=user.email,
102
- role_id=user.role_id,
102
+ role_id=user.role_id.value if user.role_id else 2,
103
)
104
logger.info(f"User: {u}")
105
session.add(u)
backend/app/connectors/models.py
+10
-9
@@ -50,17 +50,17 @@ class Connectors(SQLModel, table=True):
50
"""
51
52
id: Optional[int] = Field(default=None, primary_key=True)
53
- connector_name: str = Field()
54
- connector_type: str = Field()
55
- connector_url: str = Field()
53
+ connector_name: str = Field(max_length=256)
54
+ connector_type: str = Field(max_length=256)
55
+ connector_url: str = Field(max_length=750)
56
connector_last_updated: datetime = Field(default=datetime.utcnow())
57
- connector_username: Optional[str] = Field(default=None)
58
- connector_password: Optional[str] = Field(default=None)
59
- connector_api_key: Optional[str] = Field(default=None)
57
+ connector_username: Optional[str] = Field(default=None, max_length=256)
58
+ connector_password: Optional[str] = Field(default=None, max_length=256)
59
+ connector_api_key: Optional[str] = Field(default=None, max_length=750)
60
61
# Fields moved from ConnectorsAvailable
62
- connector_description: Optional[str] = Field(default=None)
63
- connector_supports: Optional[str] = Field(default=None)
62
+ connector_description: Optional[str] = Field(default=None, max_length=1000)
63
+ connector_supports: Optional[str] = Field(default=None, max_length=1000)
64
connector_configured: bool = Field(default=False)
65
connector_verified: bool = Field(default=False)
66
connector_accepts_host_only: bool = Field(default=False)
@@ -68,7 +68,8 @@ class Connectors(SQLModel, table=True):
68
connector_accepts_username_password: bool = Field(default=False)
69
connector_accepts_file: bool = Field(default=False)
70
connector_accepts_extra_data: bool = Field(default=False)
71
- connector_extra_data: Optional[str] = Field(default=None)
71
+ connector_extra_data: Optional[str] = Field(default=None, max_length=1000)
72
+ connector_enabled: bool = Field(default=False)
73
74
# Relationship
75
history_logs: List[ConnectorHistory] = Relationship(
backend/app/connectors/wazuh_manager/routes/rules.py
+1
-6
@@ -163,9 +163,4 @@ async def enable_wazuh_rule(
163
)
164
async def exclude_wazuh_rule() -> RuleExcludeResponse:
165
raise HTTPException(status_code=501, detail="Feature not yet ready")
166
- return RuleExcludeResponse(
167
- wazuh_rule='<group name="windows, sysmon, sysmon_event1, windows_sysmon_event1">\n<rule id="100126" level="1">\n<if_sid>100125</if_sid>\n<field name="win.eventdata.user">NT AUTHORITY\\\\SYSTEM</field>\n<field name="win.eventdata.originalFileName">Wmiprvse.exe</field>\n<field name="win.eventdata.image">C:\\\\Windows\\\\System32\\\\wbem\\\\WmiPrvSE.exe</field>\n<field name="win.eventdata.parentCommandLine">C:\\\\Windows\\\\system32\\\\svchost.exe -k DcomLaunch -p</field>\n<description>Exclusion rule for specific system processes and parent commands.</description>\n</rule>\n</group>',
168
- explanation="This rule is designed to exclude specific system processes identified by their user, original filename, image path, and parent command line. It triggers based on the presence of these attributes, which are commonly associated with legitimate system activities as defined in the payload. The rule is set to level 1 as a basic exclusion without generating an alert. It inherits from a previous rule with ID 100125. The exclusion is based on matching the exact strings for the user, original file name, process image, and parent command line execution parameters.",
169
- message="This is a test",
170
- success=True,
171
- )
166
+ return RuleExcludeResponse(success=False, message="Feature not yet ready")
backend/app/db/all_models.py
+3
-3
@@ -1,8 +1,8 @@
1
# all_models.py
2
from app.auth.models.users import User
3
from app.connectors.models import Connectors
4
-from app.connectors.sublime.models.alerts import SublimeAlerts
5
-from app.connectors.wazuh_manager.models.rules import DisabledRule
4
+# from app.connectors.sublime.models.alerts import SublimeAlerts
5
+# from app.connectors.wazuh_manager.models.rules import DisabledRule
6
from app.db.universal_models import Agents
7
from app.db.universal_models import Customers
8
from app.db.universal_models import CustomersMeta
@@ -13,7 +13,7 @@ from app.integrations.alert_creation_settings.models.alert_creation_settings imp
13
from app.integrations.models.customer_integration_settings import CustomerIntegrations
14
from app.schedulers.models.scheduler import JobMetadata
15
from app.integrations.monitoring_alert.models.monitoring_alert import MonitoringAlerts
16
-from app.integrations.sap_siem.models.sap_siem import SapSiemMultipleLogins
16
+# from app.integrations.sap_siem.models.sap_siem import SapSiemMultipleLogins
17
from app.customer_provisioning.models.default_settings import (
18
CustomerProvisioningDefaultSettings,
19
)
backend/app/db/db_populate.py
+76
-55
@@ -293,20 +293,29 @@ async def add_available_integrations_if_not_exist(session: AsyncSession):
293
available_integrations_list = get_available_integrations_list()
294
295
for available_integration_data in available_integrations_list:
296
- query = select(AvailableIntegrations).where(
297
- AvailableIntegrations.integration_name == available_integration_data["integration_name"],
298
- )
299
- result = await session.execute(query)
300
- existing_available_integration = result.scalars().first()
301
-
302
- if existing_available_integration is None:
303
- new_available_integration = AvailableIntegrations(
304
- **available_integration_data,
305
- )
306
- session.add(new_available_integration)
307
- logger.info(
308
- f"Added new available integration: {available_integration_data['integration_name']}",
296
+ try:
297
+ query = select(AvailableIntegrations).where(
298
+ AvailableIntegrations.integration_name == available_integration_data["integration_name"],
299
)
300
+ result = await session.execute(query)
301
+ existing_available_integration = result.scalars().first()
302
+
303
+ if existing_available_integration is None:
304
+ new_available_integration = AvailableIntegrations(
305
+ **available_integration_data,
306
+ )
307
+ logger.info(f"New available integration: {available_integration_data}")
308
+ session.add(new_available_integration)
309
+ logger.info(
310
+ f"Added new available integration: {available_integration_data['integration_name']}",
311
+ )
312
+ except Exception as e:
313
+ logger.error(f"Error adding available integration: {e}")
314
+ await session.rollback()
315
+ raise e
316
+ await session.commit()
317
+ # Close the session
318
+ await session.close()
319
320
321
def load_available_integrations_auth_keys(
@@ -368,24 +377,30 @@ async def get_available_integrations_auth_keys_list(session: AsyncSession):
377
("CarbonBlack", "ORGANIZATION_KEY"),
378
# ... Add more available integrations auth keys as needed ...
379
]
371
-
372
- for integration_name, auth_key_name in available_integrations:
373
- query = select(AvailableIntegrations.id).where(
374
- AvailableIntegrations.integration_name == integration_name,
375
- )
376
- result = await session.execute(query)
377
- integration_id = result.scalars().first()
378
-
379
- if integration_id:
380
- available_integrations_auth_keys.append(
381
- load_available_integrations_auth_keys(
382
- integration_id,
383
- integration_name,
384
- auth_key_name,
385
- ),
380
+ logger.info("Getting available integrations auth keys.")
381
+ try:
382
+ for integration_name, auth_key_name in available_integrations:
383
+ query = select(AvailableIntegrations.id).where(
384
+ AvailableIntegrations.integration_name == integration_name,
385
)
386
+ result = await session.execute(query)
387
+ integration_id = result.scalars().first()
388
+ logger.info(f"Integration ID for {integration_name}: {integration_id}")
389
+ if integration_id:
390
+ logger.info(f"Found integration ID for {integration_name}: {integration_id}")
391
+ available_integrations_auth_keys.append(
392
+ load_available_integrations_auth_keys(
393
+ integration_id,
394
+ integration_name,
395
+ auth_key_name,
396
+ ),
397
+ )
398
388
- return available_integrations_auth_keys
399
+ return available_integrations_auth_keys
400
+ except Exception as e:
401
+ logger.error(f"Error getting available integrations auth keys: {e}")
402
+ await session.rollback()
403
+ raise e
404
405
406
async def add_available_integrations_auth_keys_if_not_exist(session: AsyncSession):
@@ -398,33 +413,39 @@ async def add_available_integrations_auth_keys_if_not_exist(session: AsyncSessio
413
Returns:
414
None
415
"""
416
+ logger.info("Checking for existence of available integrations auth keys.")
417
available_integrations_auth_keys_list = await get_available_integrations_auth_keys_list(session=session)
402
-
418
+ logger.info("Adding available integrations auth keys to the database.")
419
for available_integration_auth_keys_data in available_integrations_auth_keys_list:
404
- query = select(AvailableIntegrations).where(
405
- AvailableIntegrations.integration_name == available_integration_auth_keys_data["integration_name"],
406
- )
407
- result = await session.execute(query)
408
- existing_integration = result.scalars().first()
409
-
410
- if existing_integration:
411
- available_integration_auth_keys_data["integration_id"] = existing_integration.id
412
- auth_key_query = select(AvailableIntegrationsAuthKeys).where(
413
- and_(
414
- AvailableIntegrationsAuthKeys.integration_id == existing_integration.id,
415
- AvailableIntegrationsAuthKeys.auth_key_name == available_integration_auth_keys_data["auth_key_name"],
416
- ),
420
+ try:
421
+ query = select(AvailableIntegrations).where(
422
+ AvailableIntegrations.integration_name == available_integration_auth_keys_data["integration_name"],
423
)
418
- auth_key_result = await session.execute(auth_key_query)
419
- existing_auth_key = auth_key_result.scalars().first()
420
-
421
- if existing_auth_key is None:
422
- new_auth_key = AvailableIntegrationsAuthKeys(
423
- **available_integration_auth_keys_data,
424
- )
425
- session.add(new_auth_key)
426
- logger.info(
427
- f"Added new available integration auth keys: "
428
- f"{available_integration_auth_keys_data['auth_key_name']} for "
429
- f"{available_integration_auth_keys_data['integration_name']}",
424
+ result = await session.execute(query)
425
+ existing_integration = result.scalars().first()
426
+
427
+ if existing_integration:
428
+ available_integration_auth_keys_data["integration_id"] = existing_integration.id
429
+ auth_key_query = select(AvailableIntegrationsAuthKeys).where(
430
+ and_(
431
+ AvailableIntegrationsAuthKeys.integration_id == existing_integration.id,
432
+ AvailableIntegrationsAuthKeys.auth_key_name == available_integration_auth_keys_data["auth_key_name"],
433
+ ),
434
)
435
+ auth_key_result = await session.execute(auth_key_query)
436
+ existing_auth_key = auth_key_result.scalars().first()
437
+
438
+ if existing_auth_key is None:
439
+ new_auth_key = AvailableIntegrationsAuthKeys(
440
+ **available_integration_auth_keys_data,
441
+ )
442
+ session.add(new_auth_key)
443
+ logger.info(
444
+ f"Added new available integration auth keys: "
445
+ f"{available_integration_auth_keys_data['auth_key_name']} for "
446
+ f"{available_integration_auth_keys_data['integration_name']}",
447
+ )
448
+ except Exception as e:
449
+ logger.error(f"Error adding available integration auth keys: {e}")
450
+ raise e
451
+ await session.commit()
backend/app/db/db_session.py
+161
-64
@@ -1,58 +1,182 @@
1
-# ! Old Testing without Async
2
-from sqlmodel import Session
3
-from sqlmodel import create_engine
4
-
5
-from settings import SQLALCHEMY_DATABASE_URI
6
-
7
-engine = create_engine(
8
- SQLALCHEMY_DATABASE_URI,
9
- connect_args={"check_same_thread": False},
10
-)
11
-session = "placeholder"
1
+# # ! Old Testing without Async
2
+# from sqlmodel import Session
3
+# from sqlmodel import create_engine
4
+
5
+# from settings import SQLALCHEMY_DATABASE_URI
6
+
7
+# engine = create_engine(
8
+# SQLALCHEMY_DATABASE_URI,
9
+# connect_args={"check_same_thread": False},
10
+# )
11
+# session = "placeholder"
12
+
13
+# from contextlib import asynccontextmanager
14
+# from contextlib import contextmanager
15
+
16
+# from loguru import logger
17
+# from sqlalchemy import create_engine
18
+# from sqlalchemy.ext.asyncio import AsyncSession
19
+# from sqlalchemy.ext.asyncio import create_async_engine
20
+# from sqlalchemy.orm import sessionmaker
21
+
22
+# from settings import SQLALCHEMY_DATABASE_URI
23
+
24
+# # create async engine for SQLite using aiosqlite
25
+# async_engine = create_async_engine(SQLALCHEMY_DATABASE_URI, echo=False)
26
+# sync_engine = create_engine(
27
+# SQLALCHEMY_DATABASE_URI.replace("+aiosqlite", ""),
28
+# echo=False,
29
+# )
30
+
31
+# # create a configured "AsyncSession" class
32
+# AsyncSessionLocal = sessionmaker(
33
+# bind=async_engine,
34
+# class_=AsyncSession,
35
+# expire_on_commit=False,
36
+# )
37
+# SyncSessionLocal = sessionmaker(
38
+# bind=sync_engine,
39
+# class_=Session,
40
+# expire_on_commit=False,
41
+# )
42
+
43
+
44
+# @asynccontextmanager
45
+# async def get_db_session():
46
+# """
47
+# Context manager that provides an asynchronous database session.
48
+
49
+# Yields:
50
+# session: An asynchronous database session.
51
+
52
+# Raises:
53
+# Exception: If an error occurs during the database session.
54
+
55
+# """
56
+# async with AsyncSessionLocal() as session:
57
+# logger.info("DB session created")
58
+# try:
59
+# yield session
60
+# except Exception as e:
61
+# logger.error(f"Error during DB session: {e}")
62
+# await session.rollback()
63
+# raise e
64
+# finally:
65
+# logger.info("Closing DB session")
66
+# await session.close()
67
+
68
+
69
+# # Synchronous context manager to get DB session for each request
70
+# @contextmanager
71
+# def get_sync_db_session():
72
+# """
73
+# Context manager that provides a synchronous database session.
74
+
75
+# Yields:
76
+# SyncSessionLocal: The synchronous database session.
77
+
78
+# Raises:
79
+# Exception: If an error occurs during the session.
80
+# """
81
+# session = SyncSessionLocal()
82
+# logger.info("Sync DB session created")
83
+# try:
84
+# yield session
85
+# except Exception as e:
86
+# logger.error(f"Error during sync DB session: {e}")
87
+# session.rollback()
88
+# raise e
89
+# finally:
90
+# logger.info("Closing sync DB session")
91
+# session.close()
92
+
93
+
94
+# @asynccontextmanager
95
+# async def get_session():
96
+# """
97
+# Context manager that provides an async session object.
98
+
99
+# Usage:
100
+# async with get_session() as session:
101
+# # Use the session object here
102
+# """
103
+# async with get_db_session() as session:
104
+# yield session
105
+
106
+
107
+# async def get_db():
108
+# """
109
+# A coroutine function that returns an asynchronous context manager for a database session.
110
+
111
+# Usage:
112
+# async with get_db() as session:
113
+# # Use the session object to interact with the database
114
+
115
+# Returns:
116
+# An asynchronous context manager that yields a database session object.
117
+# """
118
+# async with get_session() as session:
119
+# yield session
120
+
121
+
122
+# ! NEW WITH MYSQL ! #
123
124
from contextlib import asynccontextmanager
125
from contextlib import contextmanager
126
127
+# from settings import SQLALCHEMY_DATABASE_URI
128
+from pathlib import Path
129
+
130
+from environs import Env
131
from loguru import logger
17
-from sqlalchemy import create_engine
132
from sqlalchemy.ext.asyncio import AsyncSession
133
from sqlalchemy.ext.asyncio import create_async_engine
134
from sqlalchemy.orm import sessionmaker
135
+from sqlmodel import Session
136
+from sqlmodel import create_engine
137
+
138
+env = Env()
139
+env.read_env(Path(__file__).parent.parent / ".env")
140
+# env.read_env(Path(__file__).parent.parent.parent / "docker-env" / ".env")
141
+logger.info(f"Loading environment from {Path(__file__).parent.parent.parent.parent / '.env'}")
142
+
143
+db_user = env.str("MYSQL_USER", default="copilot")
144
+db_password = env.str("MYSQL_PASSWORD")
145
+db_root_password = env.str("MYSQL_ROOT_PASSWORD")
146
+db_url = env.str("MYSQL_URL", default="copilot-mysql")
147
+
148
+logger.info(f"DB User: {db_user} and password: {db_password}")
149
+
150
+# Update the SQLALCHEMY_DATABASE_URI to a MySQL compatible one in settings.py
151
+# For this example, let's assume it has been updated. copilot-mysql
152
+SQLALCHEMY_DATABASE_URI_NO_DB = f"mysql+pymysql://root:{db_root_password}@{db_url}"
153
+SQLALCHEMY_DATABASE_URI = f"mysql+aiomysql://{db_user}:{db_password}@{db_url}/copilot"
154
+
155
22
-from settings import SQLALCHEMY_DATABASE_URI
156
+session = "placeholder"
157
24
-# create async engine for SQLite using aiosqlite
25
-async_engine = create_async_engine(SQLALCHEMY_DATABASE_URI, echo=False)
158
+# Create async engine for MySQL using aiomysql
159
+async_engine = create_async_engine(
160
+ SQLALCHEMY_DATABASE_URI,
161
+ echo=False,
162
+ # Additional MySQL-specific options can be set here if needed
163
+)
164
+# If you still need sync sessions for some operations, set it appropriately
165
+# This would typically require a different sync driver since SQLAlchemy doesn't use aiomysql for sync operations
166
+# ! THIS IS USED BY THE SCHEDULER ! #
167
sync_engine = create_engine(
27
- SQLALCHEMY_DATABASE_URI.replace("+aiosqlite", ""),
168
+ SQLALCHEMY_DATABASE_URI.replace("+aiomysql", "+pymysql"),
169
echo=False,
170
+ # Additional MySQL-specific options can be set here if needed
171
)
172
31
-# create a configured "AsyncSession" class
32
-AsyncSessionLocal = sessionmaker(
33
- bind=async_engine,
34
- class_=AsyncSession,
35
- expire_on_commit=False,
36
-)
37
-SyncSessionLocal = sessionmaker(
38
- bind=sync_engine,
39
- class_=Session,
40
- expire_on_commit=False,
41
-)
173
+# Create a configured "AsyncSession" class
174
+AsyncSessionLocal = sessionmaker(bind=async_engine, class_=AsyncSession, expire_on_commit=False)
175
+SyncSessionLocal = sessionmaker(bind=sync_engine, class_=Session, expire_on_commit=False)
176
177
178
@asynccontextmanager
179
async def get_db_session():
46
- """
47
- Context manager that provides an asynchronous database session.
48
-
49
- Yields:
50
- session: An asynchronous database session.
51
-
52
- Raises:
53
- Exception: If an error occurs during the database session.
54
-
55
- """
180
async with AsyncSessionLocal() as session:
181
logger.info("DB session created")
182
try:
@@ -66,18 +190,8 @@ async def get_db_session():
190
await session.close()
191
192
69
-# Synchronous context manager to get DB session for each request
193
@contextmanager
194
def get_sync_db_session():
72
- """
73
- Context manager that provides a synchronous database session.
74
-
75
- Yields:
76
- SyncSessionLocal: The synchronous database session.
77
-
78
- Raises:
79
- Exception: If an error occurs during the session.
80
- """
195
session = SyncSessionLocal()
196
logger.info("Sync DB session created")
197
try:
@@ -93,27 +207,10 @@ def get_sync_db_session():
207
208
@asynccontextmanager
209
async def get_session():
96
- """
97
- Context manager that provides an async session object.
98
-
99
- Usage:
100
- async with get_session() as session:
101
- # Use the session object here
102
- """
210
async with get_db_session() as session:
211
yield session
212
213
214
async def get_db():
108
- """
109
- A coroutine function that returns an asynchronous context manager for a database session.
110
-
111
- Usage:
112
- async with get_db() as session:
113
- # Use the session object to interact with the database
114
-
115
- Returns:
116
- An asynchronous context manager that yields a database session object.
117
- """
215
async with get_session() as session:
216
yield session
backend/app/db/db_setup.py
+114
-2
@@ -1,11 +1,17 @@
1
+import os
2
+
3
from loguru import logger
4
+from sqlalchemy import create_engine
5
from sqlalchemy import text
6
from sqlalchemy.exc import OperationalError
7
+from sqlalchemy.exc import SQLAlchemyError
8
from sqlalchemy.ext.asyncio import AsyncSession
9
10
# ! New with Async
11
from sqlmodel import SQLModel
12
13
+from alembic import command
14
+from alembic.config import Config
15
from app.auth.services.universal import create_admin_user
16
from app.auth.services.universal import create_scheduler_user
17
from app.auth.services.universal import remove_scheduler_user
@@ -13,6 +19,107 @@ from app.db.db_populate import add_available_integrations_auth_keys_if_not_exist
19
from app.db.db_populate import add_available_integrations_if_not_exist
20
from app.db.db_populate import add_connectors_if_not_exist
21
from app.db.db_populate import add_roles_if_not_exist
22
+from app.db.db_session import SQLALCHEMY_DATABASE_URI
23
+from app.db.db_session import db_password
24
+
25
+
26
+async def create_database_if_not_exists(db_url: str, db_name: str):
27
+ """
28
+ Create a database if it does not already exist.
29
+
30
+ Args:
31
+ db_url (str): Database URL to connect to MySQL server (without database part).
32
+ db_name (str): The name of the database to create.
33
+ """
34
+ engine = create_engine(db_url)
35
+ conn = engine.connect()
36
+ try:
37
+ # Check if database exists
38
+ conn.execute("commit")
39
+ exists = conn.execute(text(f"SHOW DATABASES LIKE '{db_name}';")).fetchone()
40
+ if not exists:
41
+ # Create database if it does not exist
42
+ conn.execute("commit")
43
+ conn.execute(text(f"CREATE DATABASE {db_name};"))
44
+ logger.info(f"Database '{db_name}' created successfully.")
45
+ else:
46
+ logger.info(f"Database '{db_name}' already exists.")
47
+ except SQLAlchemyError as e:
48
+ print(f"An error occurred: {e}")
49
+ finally:
50
+ conn.close()
51
+ engine.dispose()
52
+
53
+
54
+async def create_copilot_user_if_not_exists(db_url: str, db_user_name: str):
55
+ """
56
+ Create a user if it does not already exist.
57
+
58
+ Args:
59
+ db_url (str): Database URL to connect to MySQL server (without database part).
60
+ db_user_name (str): The name of the user to create.
61
+ """
62
+ db_name = "copilot"
63
+ engine = create_engine(db_url)
64
+ conn = engine.connect()
65
+ try:
66
+ # Check if user exists
67
+ conn.execute("commit")
68
+ exists = conn.execute(text(f"SELECT * FROM mysql.user WHERE user = '{db_user_name}';")).fetchone()
69
+ if not exists:
70
+ # Create user if it does not exist
71
+ conn.execute("commit")
72
+ conn.execute(text(f"CREATE USER '{db_user_name}'@'%' IDENTIFIED BY '{db_password}';"))
73
+ logger.info(f"User '{db_user_name}' created successfully with password '{db_password}'.")
74
+ conn.execute(text(f"GRANT ALL PRIVILEGES ON {db_name}.* TO '{db_user_name}'@'%';"))
75
+ logger.info(f"User '{db_user_name}' created successfully and granted all privileges to the '{db_name}' database.")
76
+ else:
77
+ logger.info(f"User '{db_user_name}' already exists.")
78
+ except SQLAlchemyError as e:
79
+ logger.info(f"An error occurred: {e}")
80
+
81
+
82
+def apply_migrations():
83
+ """
84
+ Applies Alembic migrations to ensure the database schema is up to date.
85
+ """
86
+ logger.info("Applying migrations")
87
+
88
+ # Navigate up three levels from db_setup.py to the backend directory, then to the alembic directory
89
+ base_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
90
+ alembic_directory = os.path.join(base_dir, "alembic")
91
+
92
+ logger.info(f"base_dir: {base_dir}")
93
+ logger.info(f"Alembic directory: {alembic_directory}")
94
+
95
+ alembic_cfg = Config(os.path.join(alembic_directory, "alembic.ini"))
96
+ alembic_cfg.set_main_option("sqlalchemy.url", SQLALCHEMY_DATABASE_URI.replace("+aiomysql", "+pymysql"))
97
+ alembic_cfg.set_main_option("script_location", alembic_directory)
98
+
99
+ # Apply migrations to the latest revision
100
+ try:
101
+ command.upgrade(alembic_cfg, "head")
102
+ except OperationalError as e:
103
+ logger.error(f"Error applying migrations: {e}")
104
+ raise e
105
+
106
+
107
+async def add_connectors(async_engine):
108
+ """
109
+ Adds connectors to the database.
110
+
111
+ Args:
112
+ async_engine (AsyncEngine): The async engine used to connect to the database.
113
+
114
+ Returns:
115
+ None
116
+ """
117
+ logger.info("Adding connectors")
118
+ async with AsyncSession(
119
+ async_engine,
120
+ ) as session: # Create an AsyncSession, not just a connection
121
+ async with session.begin(): # Start a transaction
122
+ await add_connectors_if_not_exist(session)
123
124
125
async def create_tables(async_engine):
@@ -99,10 +206,15 @@ async def create_available_integrations(async_engine):
206
async with AsyncSession(
207
async_engine,
208
) as session: # Create an AsyncSession, not just a connection
102
- async with session.begin(): # Start a transaction
209
+ try:
210
await add_available_integrations_if_not_exist(session)
211
await add_available_integrations_auth_keys_if_not_exist(session)
105
- await session.commit()
212
+ except Exception as e:
213
+ logger.error(f"Error creating available integrations: {e}")
214
+ await session.rollback() # Explicit rollback on error
215
+ raise # Re-raise the exception to handle it further up the call stack
216
+ else:
217
+ await session.commit() # Explicit commit if all operations are successful
218
219
220
async def ensure_admin_user(async_engine):
backend/app/db/universal_models.py
+24
-10
@@ -1,6 +1,9 @@
1
from datetime import datetime
2
from typing import Optional
3
4
+from sqlalchemy import Column
5
+from sqlalchemy import Float
6
+from sqlalchemy import LargeBinary
7
from sqlmodel import Field
8
from sqlmodel import Relationship
9
from sqlmodel import SQLModel
@@ -8,7 +11,7 @@ from sqlmodel import SQLModel
11
12
class Customers(SQLModel, table=True):
13
id: Optional[int] = Field(primary_key=True)
11
- customer_code: str = Field(max_length=50, nullable=False)
14
+ customer_code: str = Field(sa_column_kwargs={"index": True}, max_length=50, nullable=False)
15
parent_customer_code: Optional[str] = Field(max_length=11)
16
customer_name: str = Field(max_length=50, nullable=False)
17
contact_last_name: Optional[str] = Field(max_length=50)
@@ -83,7 +86,7 @@ class CustomersMeta(SQLModel, table=True):
86
87
class Agents(SQLModel, table=True):
88
id: Optional[int] = Field(primary_key=True)
86
- agent_id: str = Field(index=True)
89
+ agent_id: str = Field(index=True, max_length=256)
90
ip_address: str = Field(max_length=256)
91
os: str = Field(max_length=256)
92
hostname: str = Field(max_length=256)
@@ -95,7 +98,7 @@ class Agents(SQLModel, table=True):
98
wazuh_agent_version: str = Field(max_length=256)
99
wazuh_agent_status: str = Field("not found", max_length=256)
100
velociraptor_agent_version: str = Field(max_length=256)
98
- customer_code: Optional[str] = Field(foreign_key="customers.customer_code")
101
+ customer_code: Optional[str] = Field(foreign_key="customers.customer_code", max_length=256)
102
quarantined: bool = Field(default=False)
103
104
customer: Optional[Customers] = Relationship(back_populates="agents")
@@ -118,8 +121,10 @@ class Agents(SQLModel, table=True):
121
wazuh_agent_version=wazuh_agent.wazuh_agent_version,
122
wazuh_agent_status=wazuh_agent.wazuh_agent_status if wazuh_agent.wazuh_agent_status else "not found",
123
velociraptor_id=velociraptor_agent.client_id if velociraptor_agent.client_id else "n/a",
121
- velociraptor_last_seen=velociraptor_agent.client_last_seen_as_datetime,
122
- velociraptor_agent_version=velociraptor_agent.client_version,
124
+ velociraptor_last_seen=velociraptor_agent.client_last_seen_as_datetime
125
+ if velociraptor_agent.client_last_seen_as_datetime
126
+ else "1970-01-01T00:00:00+00:00",
127
+ velociraptor_agent_version=velociraptor_agent.client_version if velociraptor_agent.client_version else "n/a",
128
customer_code=customer_code,
129
)
130
@@ -150,13 +155,13 @@ class LogEntry(SQLModel, table=True):
155
__tablename__ = "log_entries"
156
id: Optional[int] = Field(primary_key=True)
157
timestamp: datetime = Field(default=datetime.utcnow())
153
- event_type: str
158
+ event_type: str = Field(default="Info", max_length=256)
159
user_id: int = Field(default=None, nullable=True)
155
- route: str
156
- method: str
160
+ route: str = Field(default=None, nullable=True, max_length=256)
161
+ method: str = Field(default=None, nullable=True, max_length=256)
162
status_code: int
158
- message: str
159
- additional_info: str = Field(default=None, nullable=True)
163
+ message: str = Field(default=None, nullable=True, max_length=5024)
164
+ additional_info: str = Field(default=None, nullable=True, max_length=5024)
165
166
167
class License(SQLModel, table=True):
@@ -166,3 +171,12 @@ class License(SQLModel, table=True):
171
customer_name: str = Field(max_length=1024)
172
customer_email: str = Field(max_length=1024)
173
company_name: str = Field(max_length=1024)
174
+
175
+
176
+class SchedulerJob(SQLModel, table=True):
177
+ id: str = Field(default=None, primary_key=True, nullable=False, max_length=255)
178
+ next_run_time: float = Field(sa_column=Column(Float(), index=True))
179
+ job_state: bytes = Field(sa_column=Column(LargeBinary(), nullable=False))
180
+
181
+ def __repr__(self):
182
+ return f"<SchedulerJob(id={self.id}, next_run_time={self.next_run_time})>"
backend/app/integrations/modules/routes/carbonblack.py
+2
-2
@@ -70,10 +70,10 @@ async def get_collect_carbonblack_data(carbonblack_request, session, auth_keys):
70
@module_carbonblack_router.post(
71
"",
72
response_model=InvokeCarbonBlackResponse,
73
- description="Invoke the Huntress module.",
73
+ description="Invoke the CarbonBlack module.",
74
)
75
async def collect_carbonblack_route(carbonblack_request: InvokeCarbonBlackRequest, session: AsyncSession = Depends(get_db)):
76
- """Pull down Huntress Events."""
76
+ """Pull down CarbonBlack Events."""
77
try:
78
customer_integration_response = await get_customer_integration_response(
79
carbonblack_request.customer_code,
backend/app/schedulers/routes/scheduler.py
+42
-8
@@ -1,7 +1,9 @@
1
+import asyncio
2
from typing import Optional
3
4
from fastapi import APIRouter
5
from fastapi import Depends
6
+from fastapi import HTTPException
7
from loguru import logger
8
from sqlalchemy.ext.asyncio import AsyncSession
9
from sqlalchemy.future import select
@@ -9,15 +11,16 @@ from sqlalchemy.future import select
11
from app.db.db_session import get_db
12
from app.schedulers.models.scheduler import JobMetadata
13
from app.schedulers.scheduler import get_function_by_name
14
+from app.schedulers.scheduler import get_scheduler_instance
15
from app.schedulers.scheduler import init_scheduler
16
from app.schedulers.schema.scheduler import JobsResponse
17
18
scheduler_router = APIRouter()
19
20
18
-def get_scheduler():
21
+async def get_scheduler():
22
# Singleton pattern or reference to existing instance
20
- return init_scheduler()
23
+ return await init_scheduler()
24
25
26
async def find_job_by_id(scheduler, job_id):
@@ -88,7 +91,7 @@ async def get_all_jobs(session: AsyncSession = Depends(get_db)) -> JobsResponse:
91
JobsResponse: The response containing the list of jobs.
92
93
"""
91
- scheduler = get_scheduler()
94
+ scheduler = await get_scheduler_instance()
95
jobs = scheduler.get_jobs()
96
apscheduler_jobs = []
97
for job in jobs:
@@ -133,7 +136,7 @@ async def add_job(
136
Returns:
137
dict: A dictionary containing the success status and a message.
138
"""
136
- scheduler = get_scheduler()
139
+ scheduler = await get_scheduler_instance()
140
job_function = get_function_by_name(function_name)
141
scheduler.add_job(
142
job_function,
@@ -155,6 +158,37 @@ async def add_job(
158
return {"success": True, "message": "Job added successfully"}
159
160
161
+@scheduler_router.post("/jobs/run/{job_id}", description="Run a job")
162
+async def run_job_manually(job_id: str, session: AsyncSession = Depends(get_db)):
163
+ """
164
+ Manually triggers a scheduled job for immediate execution.
165
+
166
+ Args:
167
+ job_id (str): The identifier of the job to run.
168
+ session (AsyncSession): The database session dependency.
169
+
170
+ Returns:
171
+ A JSON response with the result of the operation.
172
+ """
173
+ scheduler = await get_scheduler_instance() # Make sure your scheduler is properly initialized
174
+ job = scheduler.get_job(job_id)
175
+
176
+ if job is None:
177
+ raise HTTPException(status_code=404, detail="Job not found")
178
+
179
+ try:
180
+ # Retrieve the function associated with the job and run it
181
+ job_function = get_function_by_name(job.name) # Ensure this function maps job names to function objects
182
+ if asyncio.iscoroutinefunction(job_function):
183
+ result = await job_function() # Execute the function if it's async
184
+ else:
185
+ result = job_function() # Execute synchronously if not an async function
186
+
187
+ return {"success": True, "message": "Job executed successfully", "result": result}
188
+ except Exception as e:
189
+ raise HTTPException(status_code=500, detail=str(e))
190
+
191
+
192
@scheduler_router.post("/start/{job_id}", description="Start a job")
193
async def start_job(job_id: str):
194
"""
@@ -168,7 +202,7 @@ async def start_job(job_id: str):
202
- If the job is found and successfully started, the success status is True and the message is "Job started successfully".
203
- If the job is not found, the success status is False and the message is "Job not found".
204
"""
171
- scheduler = get_scheduler()
205
+ scheduler = await get_scheduler_instance()
206
job = await find_job_by_id(scheduler, job_id)
207
if job:
208
job.resume()
@@ -191,7 +225,7 @@ async def pause_job(job_id: str):
225
- If the job is paused successfully, the success status is True and the message is "Job paused successfully".
226
- If the job is not found, the success status is False and the message is "Job not found".
227
"""
194
- scheduler = get_scheduler()
228
+ scheduler = await get_scheduler_instance()
229
job = await find_job_by_id(scheduler, job_id)
230
if job:
231
job.pause()
@@ -225,7 +259,7 @@ async def update_job(
259
"message": "Job updated successfully"
260
}
261
"""
228
- scheduler = get_scheduler()
262
+ scheduler = await get_scheduler_instance()
263
job = await find_job_by_id(scheduler, job_id)
264
if job:
265
job.reschedule(trigger="interval", minutes=time_interval)
@@ -263,7 +297,7 @@ async def delete_job(job_id: str, session: AsyncSession = Depends(get_db)):
297
Returns:
298
dict: A dictionary containing the success status and a message.
299
"""
266
- scheduler = get_scheduler()
300
+ scheduler = await get_scheduler_instance()
301
job = await find_job_by_id(scheduler, job_id)
302
if job:
303
scheduler.remove_job(job_id)
backend/app/schedulers/scheduler.py
+124
-39
@@ -1,8 +1,15 @@
1
+import asyncio
2
+
3
+from apscheduler.events import EVENT_JOB_ERROR
4
+from apscheduler.events import EVENT_JOB_MISSED
5
+from apscheduler.executors.asyncio import AsyncIOExecutor
6
from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore
7
from apscheduler.schedulers.asyncio import AsyncIOScheduler
8
from loguru import logger
9
+from sqlalchemy.ext.asyncio import AsyncSession
10
+from sqlalchemy.future import select
11
5
-from app.db.db_session import SyncSessionLocal
12
+from app.db.db_session import async_engine
13
from app.db.db_session import sync_engine
14
from app.schedulers.models.scheduler import CreateSchedulerRequest
15
from app.schedulers.models.scheduler import JobMetadata
@@ -49,36 +56,75 @@ from app.schedulers.services.monitoring_alert import invoke_suricata_monitoring_
56
from app.schedulers.services.monitoring_alert import invoke_wazuh_monitoring_alert
57
58
52
-def init_scheduler():
53
- """
54
- Initializes and configures the scheduler.
55
- """
56
- scheduler = AsyncIOScheduler()
57
- jobstores = {"default": SQLAlchemyJobStore(engine=sync_engine)}
58
- scheduler.configure(jobstores=jobstores)
59
+def scheduler_listener(event):
60
+ if event.exception:
61
+ logger.error(f"Job {event.job_id} crashed: {event.exception}")
62
+ else:
63
+ logger.info(
64
+ f"Job {event.job_id} that was scheduled to run at {event.scheduled_run_time}, missed its run time by {event.scheduled_run_time - event.scheduled_run_time}",
65
+ )
66
60
- initialize_job_metadata()
61
- schedule_enabled_jobs(scheduler)
67
63
- if not scheduler.running:
64
- scheduler.start()
68
+# Global variable to hold the scheduler instance
69
+scheduler_instance = None
70
+
71
+
72
+async def init_scheduler():
73
+ global scheduler_instance
74
+ if scheduler_instance is not None:
75
+ logger.info("Returning existing scheduler instance.")
76
+ return scheduler_instance
77
+
78
+ logger.info("Initializing new scheduler...")
79
+ try:
80
+ jobstores = {"default": SQLAlchemyJobStore(engine=sync_engine, tablename="schedulerjob")}
81
+ executors = {"default": AsyncIOExecutor()} # This executor can run asyncio coroutines
82
+ event_loop = asyncio.get_event_loop()
83
+ scheduler_instance = AsyncIOScheduler(event_loop=event_loop)
84
+ scheduler_instance.add_listener(scheduler_listener, EVENT_JOB_MISSED | EVENT_JOB_ERROR)
85
+ scheduler_instance.configure(jobstores=jobstores, executors=executors)
86
+ await initialize_job_metadata()
87
+ logger.info("Scheduling enabled jobs...")
88
+ await schedule_enabled_jobs(scheduler_instance)
89
+
90
+ if not scheduler_instance.running:
91
+ logger.info("Starting scheduler...")
92
+ scheduler_instance.start()
93
+ logger.info("Scheduler started.")
94
+
95
+ except Exception as e:
96
+ logger.error(f"Error initializing scheduler: {e}")
97
+ raise
98
66
- return scheduler
99
+ return scheduler_instance
100
101
69
-def initialize_job_metadata():
102
+async def get_scheduler_instance():
103
+ """
104
+ Retrieves the current scheduler instance. Initializes one if it does not exist.
105
+ """
106
+ global scheduler_instance
107
+ if scheduler_instance is None:
108
+ return await init_scheduler()
109
+ return scheduler_instance
110
+
111
+
112
+async def initialize_job_metadata():
113
"""
114
Initializes job metadata from the database.
115
"""
73
- with SyncSessionLocal() as session:
116
+ async with AsyncSession(async_engine) as session:
117
# Implement logic to initialize or update job metadata.
118
# Example: Check and add metadata for each known job
119
known_jobs = [
77
- {"job_id": "agent_sync", "time_interval": 15, "function": agent_sync},
120
+ {"job_id": "agent_sync", "time_interval": 1, "function": agent_sync},
121
# {"job_id": "invoke_mimecast_integration", "time_interval": 5, "function": invoke_mimecast_integration}
122
]
123
for job in known_jobs:
81
- job_metadata = session.query(JobMetadata).filter_by(job_id=job["job_id"]).one_or_none()
124
+ # Create a select statement for the JobMetadata table
125
+ stmt = select(JobMetadata).where(JobMetadata.job_id == job["job_id"])
126
+ result = await session.execute(stmt)
127
+ job_metadata = result.scalars().one_or_none()
128
if not job_metadata:
129
job_metadata = JobMetadata(
130
job_id=job["job_id"],
@@ -90,25 +136,43 @@ def initialize_job_metadata():
136
else:
137
job_metadata.time_interval = job["time_interval"]
138
job_metadata.enabled = True
93
- session.commit()
139
+ await session.commit()
140
141
96
-def schedule_enabled_jobs(scheduler):
142
+async def schedule_enabled_jobs(scheduler):
143
"""
144
Schedules jobs that are enabled in the database.
145
"""
100
- with SyncSessionLocal() as session:
101
- job_metadatas = session.query(JobMetadata).filter_by(enabled=True).all()
146
+ async with AsyncSession(async_engine) as session:
147
+ stmt = select(JobMetadata).where(JobMetadata.enabled == True)
148
+ result = await session.execute(stmt)
149
+ job_metadatas = result.scalars().all()
150
+
151
for job_metadata in job_metadatas:
152
try:
153
job_function = get_function_by_name(job_metadata.job_id)
105
- scheduler.add_job(
106
- job_function,
107
- "interval",
108
- minutes=job_metadata.time_interval,
109
- id=job_metadata.job_id,
110
- replace_existing=True,
111
- )
154
+ if asyncio.iscoroutinefunction(job_function):
155
+ # Adding coroutine functions directly
156
+ scheduler.add_job(
157
+ job_function,
158
+ "interval",
159
+ minutes=job_metadata.time_interval,
160
+ id=job_metadata.job_id,
161
+ replace_existing=True,
162
+ coalesce=True,
163
+ max_instances=1,
164
+ )
165
+ else:
166
+ # Regular functions go here
167
+ scheduler.add_job(
168
+ job_function,
169
+ "interval",
170
+ minutes=job_metadata.time_interval,
171
+ id=job_metadata.job_id,
172
+ replace_existing=True,
173
+ coalesce=True,
174
+ max_instances=1,
175
+ )
176
logger.info(f"Scheduled job: {job_metadata.job_id}")
177
except ValueError as e:
178
logger.error(f"Error scheduling job: {e}")
@@ -153,18 +217,34 @@ async def add_scheduler_jobs(create_scheduler_request: CreateSchedulerRequest):
217
Args:
218
create_scheduler_request (CreateSchedulerRequest): The request object containing the job details.
219
"""
156
- scheduler = init_scheduler()
220
+ scheduler = await get_scheduler_instance()
221
logger.info(f"create_scheduler_request: {create_scheduler_request}")
222
223
job_function = get_function_by_name(create_scheduler_request.function_name)
224
161
- scheduler.add_job(
162
- job_function,
163
- "interval",
164
- minutes=create_scheduler_request.time_interval,
165
- id=create_scheduler_request.job_id,
166
- replace_existing=True,
167
- )
225
+ # Here, we use the async add_job if the job function is a coroutine
226
+ if asyncio.iscoroutinefunction(job_function):
227
+ # Adding coroutine functions directly
228
+ scheduler.add_job(
229
+ job_function,
230
+ "interval",
231
+ minutes=create_scheduler_request.time_interval,
232
+ id=create_scheduler_request.job_id,
233
+ replace_existing=True,
234
+ coalesce=True,
235
+ max_instances=1,
236
+ )
237
+ else:
238
+ # Regular functions go here
239
+ scheduler.add_job(
240
+ job_function,
241
+ "interval",
242
+ minutes=create_scheduler_request.time_interval,
243
+ id=create_scheduler_request.job_id,
244
+ replace_existing=True,
245
+ coalesce=True,
246
+ max_instances=1,
247
+ )
248
249
await add_job_metadata(create_scheduler_request)
250
@@ -179,8 +259,12 @@ async def add_job_metadata(create_scheduler_request: CreateSchedulerRequest):
259
Args:
260
create_scheduler_request (CreateSchedulerRequest): The request object containing the job details.
261
"""
182
- with SyncSessionLocal() as session:
183
- job_metadata = session.query(JobMetadata).filter_by(job_id=create_scheduler_request.job_id).one_or_none()
262
+ async with AsyncSession(async_engine) as session:
263
+ # Using SQLAlchemy 1.4+ style with select() and scalars() for fetching results
264
+ stmt = select(JobMetadata).where(JobMetadata.job_id == create_scheduler_request.job_id)
265
+ result = await session.execute(stmt)
266
+ job_metadata = result.scalars().one_or_none()
267
+
268
if not job_metadata:
269
job_metadata = JobMetadata(
270
job_id=create_scheduler_request.job_id,
@@ -192,4 +276,5 @@ async def add_job_metadata(create_scheduler_request: CreateSchedulerRequest):
276
else:
277
job_metadata.time_interval = create_scheduler_request.time_interval
278
job_metadata.enabled = True
195
- session.commit()
279
+
280
+ await session.commit()
backend/app/schedulers/services/agent_sync.py
+10
-8
@@ -1,10 +1,11 @@
1
from datetime import datetime
2
3
from dotenv import load_dotenv
4
+from loguru import logger
5
+from sqlalchemy.future import select
6
7
from app.agents.routes.agents import sync_all_agents
8
from app.db.db_session import get_db_session
7
-from app.db.db_session import get_sync_db_session
9
from app.schedulers.models.scheduler import JobMetadata
10
11
load_dotenv()
@@ -20,17 +21,18 @@ async def agent_sync():
21
If the token retrieval fails, it prints a failure message. If the job metadata for
22
'agent_sync' does not exist, it prints a message indicating the absence of the metadata.
23
"""
24
+ logger.info("Synchronizing agents via scheduler...")
25
async with get_db_session() as session:
26
await sync_all_agents(session=session)
27
26
- # Use get_sync_db_session to create and manage a synchronous session
27
- with get_sync_db_session() as session:
28
- # Synchronous ORM operations
29
- job_metadata = session.query(JobMetadata).filter_by(job_id="agent_sync").one_or_none()
28
+ stmt = select(JobMetadata).where(JobMetadata.job_id == "agent_sync")
29
+ result = await session.execute(stmt)
30
+ job_metadata = result.scalars().first()
31
+
32
if job_metadata:
33
job_metadata.last_success = datetime.utcnow()
34
session.add(job_metadata)
33
- session.commit()
35
+ await session.commit() # Asynchronously commit the transaction
36
+ logger.info("Updated job metadata with the last success timestamp.")
37
else:
35
- # Handle the case where job_metadata does not exist
36
- print("JobMetadata for 'agent_sync' not found.")
38
+ logger.warning("JobMetadata for 'agent_sync' not found.")
backend/copilot.py
+19
-8
@@ -10,14 +10,17 @@ from fastapi.middleware.cors import CORSMiddleware
10
from loguru import logger
11
12
from app.auth.utils import AuthHandler
13
+from app.db.db_session import SQLALCHEMY_DATABASE_URI_NO_DB
14
from app.db.db_session import async_engine
15
+from app.db.db_setup import add_connectors
16
+from app.db.db_setup import apply_migrations
17
from app.db.db_setup import create_available_integrations
18
+from app.db.db_setup import create_copilot_user_if_not_exists
19
+from app.db.db_setup import create_database_if_not_exists
20
from app.db.db_setup import create_roles
16
-from app.db.db_setup import create_tables
21
from app.db.db_setup import ensure_admin_user
22
from app.db.db_setup import ensure_scheduler_user
23
from app.db.db_setup import ensure_scheduler_user_removed
20
-from app.db.db_setup import update_tables
24
from app.middleware.exception_handlers import custom_http_exception_handler
25
from app.middleware.exception_handlers import validation_exception_handler
26
from app.middleware.exception_handlers import value_error_handler
@@ -57,12 +60,15 @@ from app.routers import threat_intel
60
from app.routers import velociraptor
61
from app.routers import wazuh_indexer
62
from app.routers import wazuh_manager
63
+from app.schedulers.scheduler import get_scheduler_instance
64
from app.schedulers.scheduler import init_scheduler
65
66
auth_handler = AuthHandler()
67
# Get the `SERVER_IP` from the `.env` file
68
load_dotenv()
69
server_ip = os.getenv("SERVER_IP", "localhost")
70
+environment = os.getenv("ENVIRONMENT", "PRODUCTION")
71
+
72
# Not needed for now
73
# ssl_keyfile = os.path.join(os.path.dirname(__file__), "../nginx/server.key")
74
# ssl_certfile = os.path.join(os.path.dirname(__file__), "../nginx/server.crt")
@@ -84,6 +90,7 @@ app.add_middleware(
90
91
92
################## ! Middleware LOGGING TO `log_entry` table ! ##################
93
+# Comment out logging for now, not sure I want to use it
94
app.middleware("http")(log_requests) # using the imported middleware
95
96
@@ -136,19 +143,22 @@ app.include_router(api_router)
143
144
@app.on_event("startup")
145
async def init_db():
139
- # create_tables(engine)
140
- await create_tables(async_engine)
141
- await update_tables(async_engine)
146
+ logger.info("Initializing database")
147
+ if environment == "PRODUCTION":
148
+ await create_database_if_not_exists(db_url=SQLALCHEMY_DATABASE_URI_NO_DB, db_name="copilot")
149
+ await create_copilot_user_if_not_exists(db_url=SQLALCHEMY_DATABASE_URI_NO_DB, db_user_name="copilot")
150
+ apply_migrations()
151
+ await add_connectors(async_engine)
152
await create_roles(async_engine)
153
await create_available_integrations(async_engine)
154
await ensure_admin_user(async_engine)
155
await ensure_scheduler_user(async_engine)
156
157
# Initialize the scheduler
148
- scheduler = init_scheduler()
158
+ scheduler = await init_scheduler()
159
150
- logger.info("Starting scheduler")
160
if not scheduler.running:
161
+ logger.info("Scheduler is not running, starting now...")
162
scheduler.start()
163
164
@@ -161,8 +171,9 @@ def hello():
171
async def shutdown_scheduler():
172
logger.info("Shutting down scheduler")
173
# Initialize the scheduler
164
- scheduler = init_scheduler()
174
+ scheduler = await get_scheduler_instance()
175
if scheduler.running:
176
+ logger.info("Scheduler is running, shutting down now...")
177
scheduler.shutdown()
178
179
await ensure_scheduler_user_removed(async_engine)
backend/data/api.config.yaml
new
+72
@@ -0,0 +1,72 @@
1
+ca_certificate: |
2
+ -----BEGIN CERTIFICATE-----
3
+ MIIDTDCCAjSgAwIBAgIRAKim5DSDvIpnbc7BN8WxDrAwDQYJKoZIhvcNAQELBQAw
4
+ GjEYMBYGA1UEChMPVmVsb2NpcmFwdG9yIENBMB4XDTIzMDYwMjEzMzU0MFoXDTMz
5
+ MDUzMDEzMzU0MFowGjEYMBYGA1UEChMPVmVsb2NpcmFwdG9yIENBMIIBIjANBgkq
6
+ hkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAu6aTXL1/NcNsrWbqn0SWZk4d5fbwz48+
7
+ SYcbZxM8aseCv7LlZEyA+aOQXyFFzdX+B3E+7+25+bfmEU6B0/9N5S0Qk+bkhw14
8
+ 0+Edk9uC9qEW2LDTNAH34T4Znb2ki+OjYgB78MWfKm7AR4KxM1wfgOq9VTQtF3Vi
9
+ HwieHyYkvnmwedA6McA/SxwY05XTlCOgrtenDRyDP2fRVPPbj6vVdLHb3EpjxpKP
10
+ 0rB/h1hoePaQ0l/AGZ8kWV2seCkmYkf+drbqxzHre6tbzawJjngcu2/FwW2J6yfR
11
+ Xcx8ETM7o8iAuSGPWoMAljjND2+bJRz2t6GJibL749tkge2lE5NnRQIDAQABo4GM
12
+ MIGJMA4GA1UdDwEB/wQEAwICpDAdBgNVHSUEFjAUBggrBgEFBQcDAQYIKwYBBQUH
13
+ AwIwDwYDVR0TAQH/BAUwAwEB/zAdBgNVHQ4EFgQUiXe4VSsY3g0DVOwxHGl7SDhV
14
+ zm8wKAYDVR0RBCEwH4IdVmVsb2NpcmFwdG9yX2NhLnZlbG9jaWRleC5jb20wDQYJ
15
+ KoZIhvcNAQELBQADggEBABPWrJ0TuRuJYVdvz8qEhW/ZhFC4fs0cPcPvfQBNiBW1
16
+ n/6esctopeDw6wW9A+cTf2jHqnBSNosDOcATa+JDR3tbq6qHAbso6FkZlgcmmYkN
17
+ qwcmeJMedym7UMQGSkN+PqfLx4nPGrMkDhsmacYM1IJ0mlGwnTmMqeA0/oRXNVEg
18
+ x4kgieeYa5J6K48CSyWAgwwCJ03vWJ+n3cpD+hWuVmK1tn3To05AG6gHMUeSK17N
19
+ qIz+2JyvBBBlwgTYUUmEYzgjNYKP0Crx57jvJZ8vs/vadpXdU29UOzeViyFEyRV4
20
+ JIo2Kx/jaPfLCPd9oE37KekiGBBkCgtxVp6sKHUYX3o=
21
+ -----END CERTIFICATE-----
22
+client_cert: |
23
+ -----BEGIN CERTIFICATE-----
24
+ MIIDWjCCAkKgAwIBAgIQS/Lj2MAoxGZCF1if5sfkUzANBgkqhkiG9w0BAQsFADAa
25
+ MRgwFgYDVQQKEw9WZWxvY2lyYXB0b3IgQ0EwHhcNMjMwNjAyMTM0MzMyWhcNMjQw
26
+ NjAxMTM0MzMyWjA1MRUwEwYDVQQKEwxWZWxvY2lyYXB0b3IxHDAaBgNVBAMME2lu
27
+ Zm9Ac29jZm9ydHJlc3MuY28wggEiMA0GCSqGSIb3DQEBAQUAA4IBDwAwggEKAoIB
28
+ AQC63lhkXRY+x6ceGQv008OrmZuG50xORhG6iYCABtkwd3scQmMTIGVMPL6Bergr
29
+ s1F/d+jDVJnannrjwdY0yiP1yIgNFNVEr7li5IeSsJak58qMs2cGI6Yb3BajbtUf
30
+ dFXCIADF3wUJNJEoXLOdpOL4sVsrAk9tX8XErd8iSjUqNIFykMD26YTxtM8cyQgB
31
+ X9XYO0OIdMJW3TSQeQYDmS22C9v9k3wNVc3Sz2TOgxrcbbva6nA++OlYQnBfta3n
32
+ QwZybJWvbUhkxrnd8Wlsu+Lvab7DGarYvaMb7c5D/YV3jyWv7K4wuFPYDmU+yuDS
33
+ uBP+fF3YNXX7KUm9ki9VQBzZAgMBAAGjgYAwfjAOBgNVHQ8BAf8EBAMCBaAwHQYD
34
+ VR0lBBYwFAYIKwYBBQUHAwEGCCsGAQUFBwMCMAwGA1UdEwEB/wQCMAAwHwYDVR0j
35
+ BBgwFoAUiXe4VSsY3g0DVOwxHGl7SDhVzm8wHgYDVR0RBBcwFYITaW5mb0Bzb2Nm
36
+ b3J0cmVzcy5jbzANBgkqhkiG9w0BAQsFAAOCAQEAYYex+11Q1n5hFen2R+4kwHAy
37
+ 7EhhLOhVuPDqVrQ+cYbUVWqJmf+x+HPpl8W/8bfoOefCy7r09Bpi8q7yBAKH6q3J
38
+ NqkZJDByaE92UAa64M5W16ptomsA51hh4PWEsLb55IVFw0h7thSl/2daaJX1x24U
39
+ oyac+rm/99NwLctJ1zBCnOWufPcwxmMLDKkBbO/cXDvgOLafqeYdagO9ssEWSI5T
40
+ FSl10PNE3c2cslDwnUU5/HZ99xRa+wDv/l7zcoe81+vgGHzIAm14aOVSnBukB+Ht
41
+ evdf3hnkPXNPUPq/yt/iV3VD7Dlb0AMU8ofQl9niuaX/ZSLoueSAXrLaOJgFqA==
42
+ -----END CERTIFICATE-----
43
+client_private_key: |
44
+ -----BEGIN RSA PRIVATE KEY-----
45
+ MIIEpAIBAAKCAQEAut5YZF0WPsenHhkL9NPDq5mbhudMTkYRuomAgAbZMHd7HEJj
46
+ EyBlTDy+gXq4K7NRf3fow1SZ2p5648HWNMoj9ciIDRTVRK+5YuSHkrCWpOfKjLNn
47
+ BiOmG9wWo27VH3RVwiAAxd8FCTSRKFyznaTi+LFbKwJPbV/FxK3fIko1KjSBcpDA
48
+ 9umE8bTPHMkIAV/V2DtDiHTCVt00kHkGA5kttgvb/ZN8DVXN0s9kzoMa3G272upw
49
+ PvjpWEJwX7Wt50MGcmyVr21IZMa53fFpbLvi72m+wxmq2L2jG+3OQ/2Fd48lr+yu
50
+ MLhT2A5lPsrg0rgT/nxd2DV1+ylJvZIvVUAc2QIDAQABAoIBAFoBajW9GE/Yvkei
51
+ 7L1Zmi925xBNK9Wvri5YuEnyAn5zrhpoZ2v4+JGF2IRo5Xg3AJQS30vl3c0M9Efr
52
+ Pw9iJXvmwJD8bdSNhw9430vqPkTjWA35AzBTz1gv47+ITKK/1+aOn5Cu4LAUX64/
53
+ KExP9PqwAidvD26w6ILY9FaBw6W1nZy5ZfpZi4RnnKmnBxTpzO0yTBgT9nBlRhYU
54
+ BmtDq5PcO0G8THKso1N/1IT2mKTA+475OjSrnJpXb43itiNazz+u5c4WMnlGRbBY
55
+ qPlu+pqQ99Jf6TefsILRMEDBNsYiIydwqfp2CRgZWblMero8wVuDIG5fwzqWvtof
56
+ 71lDJ3kCgYEA49E9F5lenpSx2tyTNS+Y8RP16mwUkvKPuxBIT63kzaRCh+1XLeNg
57
+ TNdnkXHLEohsdLYewITsgNOrwMtMMWguSAlbyK0VgZvgESpAlrwWR6RSTSHJeLNW
58
+ 2yj7ZkUgkSooxGv3KCh4tSJ4cwzT0N8YtHwB/U/G8w4A+fv355spmgMCgYEA0fxL
59
+ 5xnHg9fzKdTWIsjQ+9byRmwhAQxkhPDrE2kq+A31xz3l/0+IKjhiIdxHrg2Qi+aN
60
+ noSnXRl0j1LdbTNPTyt2AF6u59yR3owYS7kIc22kw/OuNdDKa1hTAPd4mucwvwlf
61
+ vJp0kMqek5hPwMKHYxNunji5oSAKLeoQrJKzpPMCgYAJksdgcH97ZoA61D4TZBan
62
+ OtGAsl4C9tJ3Z+3B+2q8AYUSNTUOpplrYTnm8MM52iXEmcqdCHjvyPVUurZO9TPM
63
+ ryf+PNfEhIpb7kyciPbet9cFir/upIqn1wcJeyotL3pbFrZiJ6E662HoY8ea5WUi
64
+ YHus62dO223LE32NbAXJhQKBgQCNyJgLm/l+SWLTvPU1mXiag0ElUb9bMN3ycaWY
65
+ fqtXwD1S4bWZlT7wmw+Po2f22wvdmrfG7/T9xMFSQPZn1HxZjZPongXlYqZPqTKC
66
+ dpaBMehNswzPI4J5xrKM9YvGtBHS++Zbt8K7PUSjjfVTx0WZHTbUuKKGa9bTt7a4
67
+ f3mzBwKBgQCaL8PuI05KZ0WxzY2b4WFErMBO01iG4ueVhslCn1c1w5GNlRDRh/Gy
68
+ ocit+c0jeOBkA+JjkyFjf7KhgzPXJQ/BmS1UeaNStUy/TPY/gx3mjuMmli+i/Z7V
69
+ drZUWaGTet9/lZLvp1jYg2shnsRyaexXUV9wlrkPft5TJhHNKvAVOw==
70
+ -----END RSA PRIVATE KEY-----
71
+api_connection_string: ashvlo01.socfortress.local:8001
72
+name: info@socfortress.co
backend/requirements.txt
+3
@@ -1,8 +1,10 @@
1
aiocsv==1.2.5
2
aiofiles==23.2.1
3
aiohttp==3.8.5
4
+aiomysql==0.2.0
5
aiosignal==1.3.1
6
aiosqlite==0.19.0
7
+alembic==1.13.1
8
amqp==5.1.1
9
annotated-types==0.6.0
10
antlr4-python3-runtime==4.9.3
@@ -112,6 +114,7 @@ pycparser==2.21
114
# pydantic-settings==2.0.3
115
Pygments==2.16.1
116
PyJWT==2.8.0
117
+PyMySQL==1.1.0
118
pyparsing==3.1.1
119
pyrsistent==0.19.3
120
pytest==7.4.2
backend/wait-for-it.sh
new
+182
@@ -0,0 +1,182 @@
1
+#!/usr/bin/env bash
2
+# Use this script to test if a given TCP host/port are available
3
+
4
+WAITFORIT_cmdname=${0##*/}
5
+
6
+echoerr() { if [[ $WAITFORIT_QUIET -ne 1 ]]; then echo "$@" 1>&2; fi }
7
+
8
+usage()
9
+{
10
+ cat << USAGE >&2
11
+Usage:
12
+ $WAITFORIT_cmdname host:port [-s] [-t timeout] [-- command args]
13
+ -h HOST | --host=HOST Host or IP under test
14
+ -p PORT | --port=PORT TCP port under test
15
+ Alternatively, you specify the host and port as host:port
16
+ -s | --strict Only execute subcommand if the test succeeds
17
+ -q | --quiet Don't output any status messages
18
+ -t TIMEOUT | --timeout=TIMEOUT
19
+ Timeout in seconds, zero for no timeout
20
+ -- COMMAND ARGS Execute command with args after the test finishes
21
+USAGE
22
+ exit 1
23
+}
24
+
25
+wait_for()
26
+{
27
+ if [[ $WAITFORIT_TIMEOUT -gt 0 ]]; then
28
+ echoerr "$WAITFORIT_cmdname: waiting $WAITFORIT_TIMEOUT seconds for $WAITFORIT_HOST:$WAITFORIT_PORT"
29
+ else
30
+ echoerr "$WAITFORIT_cmdname: waiting for $WAITFORIT_HOST:$WAITFORIT_PORT without a timeout"
31
+ fi
32
+ WAITFORIT_start_ts=$(date +%s)
33
+ while :
34
+ do
35
+ if [[ $WAITFORIT_ISBUSY -eq 1 ]]; then
36
+ nc -z $WAITFORIT_HOST $WAITFORIT_PORT
37
+ WAITFORIT_result=$?
38
+ else
39
+ (echo -n > /dev/tcp/$WAITFORIT_HOST/$WAITFORIT_PORT) >/dev/null 2>&1
40
+ WAITFORIT_result=$?
41
+ fi
42
+ if [[ $WAITFORIT_result -eq 0 ]]; then
43
+ WAITFORIT_end_ts=$(date +%s)
44
+ echoerr "$WAITFORIT_cmdname: $WAITFORIT_HOST:$WAITFORIT_PORT is available after $((WAITFORIT_end_ts - WAITFORIT_start_ts)) seconds"
45
+ break
46
+ fi
47
+ sleep 1
48
+ done
49
+ return $WAITFORIT_result
50
+}
51
+
52
+wait_for_wrapper()
53
+{
54
+ # In order to support SIGINT during timeout: http://unix.stackexchange.com/a/57692
55
+ if [[ $WAITFORIT_QUIET -eq 1 ]]; then
56
+ timeout $WAITFORIT_BUSYTIMEFLAG $WAITFORIT_TIMEOUT $0 --quiet --child --host=$WAITFORIT_HOST --port=$WAITFORIT_PORT --timeout=$WAITFORIT_TIMEOUT &
57
+ else
58
+ timeout $WAITFORIT_BUSYTIMEFLAG $WAITFORIT_TIMEOUT $0 --child --host=$WAITFORIT_HOST --port=$WAITFORIT_PORT --timeout=$WAITFORIT_TIMEOUT &
59
+ fi
60
+ WAITFORIT_PID=$!
61
+ trap "kill -INT -$WAITFORIT_PID" INT
62
+ wait $WAITFORIT_PID
63
+ WAITFORIT_RESULT=$?
64
+ if [[ $WAITFORIT_RESULT -ne 0 ]]; then
65
+ echoerr "$WAITFORIT_cmdname: timeout occurred after waiting $WAITFORIT_TIMEOUT seconds for $WAITFORIT_HOST:$WAITFORIT_PORT"
66
+ fi
67
+ return $WAITFORIT_RESULT
68
+}
69
+
70
+# process arguments
71
+while [[ $# -gt 0 ]]
72
+do
73
+ case "$1" in
74
+ *:* )
75
+ WAITFORIT_hostport=(${1//:/ })
76
+ WAITFORIT_HOST=${WAITFORIT_hostport[0]}
77
+ WAITFORIT_PORT=${WAITFORIT_hostport[1]}
78
+ shift 1
79
+ ;;
80
+ --child)
81
+ WAITFORIT_CHILD=1
82
+ shift 1
83
+ ;;
84
+ -q | --quiet)
85
+ WAITFORIT_QUIET=1
86
+ shift 1
87
+ ;;
88
+ -s | --strict)
89
+ WAITFORIT_STRICT=1
90
+ shift 1
91
+ ;;
92
+ -h)
93
+ WAITFORIT_HOST="$2"
94
+ if [[ $WAITFORIT_HOST == "" ]]; then break; fi
95
+ shift 2
96
+ ;;
97
+ --host=*)
98
+ WAITFORIT_HOST="${1#*=}"
99
+ shift 1
100
+ ;;
101
+ -p)
102
+ WAITFORIT_PORT="$2"
103
+ if [[ $WAITFORIT_PORT == "" ]]; then break; fi
104
+ shift 2
105
+ ;;
106
+ --port=*)
107
+ WAITFORIT_PORT="${1#*=}"
108
+ shift 1
109
+ ;;
110
+ -t)
111
+ WAITFORIT_TIMEOUT="$2"
112
+ if [[ $WAITFORIT_TIMEOUT == "" ]]; then break; fi
113
+ shift 2
114
+ ;;
115
+ --timeout=*)
116
+ WAITFORIT_TIMEOUT="${1#*=}"
117
+ shift 1
118
+ ;;
119
+ --)
120
+ shift
121
+ WAITFORIT_CLI=("$@")
122
+ break
123
+ ;;
124
+ --help)
125
+ usage
126
+ ;;
127
+ *)
128
+ echoerr "Unknown argument: $1"
129
+ usage
130
+ ;;
131
+ esac
132
+done
133
+
134
+if [[ "$WAITFORIT_HOST" == "" || "$WAITFORIT_PORT" == "" ]]; then
135
+ echoerr "Error: you need to provide a host and port to test."
136
+ usage
137
+fi
138
+
139
+WAITFORIT_TIMEOUT=${WAITFORIT_TIMEOUT:-15}
140
+WAITFORIT_STRICT=${WAITFORIT_STRICT:-0}
141
+WAITFORIT_CHILD=${WAITFORIT_CHILD:-0}
142
+WAITFORIT_QUIET=${WAITFORIT_QUIET:-0}
143
+
144
+# Check to see if timeout is from busybox?
145
+WAITFORIT_TIMEOUT_PATH=$(type -p timeout)
146
+WAITFORIT_TIMEOUT_PATH=$(realpath $WAITFORIT_TIMEOUT_PATH 2>/dev/null || readlink -f $WAITFORIT_TIMEOUT_PATH)
147
+
148
+WAITFORIT_BUSYTIMEFLAG=""
149
+if [[ $WAITFORIT_TIMEOUT_PATH =~ "busybox" ]]; then
150
+ WAITFORIT_ISBUSY=1
151
+ # Check if busybox timeout uses -t flag
152
+ # (recent Alpine versions don't support -t anymore)
153
+ if timeout &>/dev/stdout | grep -q -e '-t '; then
154
+ WAITFORIT_BUSYTIMEFLAG="-t"
155
+ fi
156
+else
157
+ WAITFORIT_ISBUSY=0
158
+fi
159
+
160
+if [[ $WAITFORIT_CHILD -gt 0 ]]; then
161
+ wait_for
162
+ WAITFORIT_RESULT=$?
163
+ exit $WAITFORIT_RESULT
164
+else
165
+ if [[ $WAITFORIT_TIMEOUT -gt 0 ]]; then
166
+ wait_for_wrapper
167
+ WAITFORIT_RESULT=$?
168
+ else
169
+ wait_for
170
+ WAITFORIT_RESULT=$?
171
+ fi
172
+fi
173
+
174
+if [[ $WAITFORIT_CLI != "" ]]; then
175
+ if [[ $WAITFORIT_RESULT -ne 0 && $WAITFORIT_STRICT -eq 1 ]]; then
176
+ echoerr "$WAITFORIT_cmdname: strict mode, refusing to execute subprocess"
177
+ exit $WAITFORIT_RESULT
178
+ fi
179
+ exec "${WAITFORIT_CLI[@]}"
180
+else
181
+ exit $WAITFORIT_RESULT
182
+fi
docker-compose.yml
+17
@@ -11,6 +11,8 @@ services:
11
# Mount the copilot.db file to persist the database
12
- ./data/data:/opt/copilot/backend/data
13
env_file: .env
14
+ depends_on:
15
+ - copilot-mysql
16
17
copilot-frontend:
18
image: ghcr.io/socfortress/copilot-frontend:latest
@@ -20,6 +22,21 @@ services:
22
- "80:80"
23
- "443:443"
24
25
+ copilot-mysql:
26
+ image: mysql:5.7
27
+ environment:
28
+ MYSQL_ROOT_PASSWORD: ${MYSQL_ROOT_PASSWORD}
29
+ MYSQL_DATABASE: copilot
30
+ MYSQL_USER: ${MYSQL_USER}
31
+ MYSQL_PASSWORD: ${MYSQL_PASSWORD}
32
+ ports:
33
+ - "3306:3306"
34
+ volumes:
35
+ - mysql-data:/var/lib/mysql
36
+
37
+volumes:
38
+ mysql-data:
39
+
40
networks:
41
default:
42
driver: bridge