This commit is contained in:
Chris Milne 2026-05-06 12:02:26 +01:00
commit ac39206625
58 changed files with 1814 additions and 0 deletions

View file

@ -0,0 +1,7 @@
"""
Configurations for <this module>
Configurations:
- List: Description
- Configs: Description
"""

View file

@ -0,0 +1,7 @@
"""
Constants and error codes for <this module>
Constants:
- List: Description
- Consts: Description
"""

View file

@ -0,0 +1,11 @@
"""
Router dependencies for <this module>
Classes:
- List: Description
- Classes: Description
Functions:
- List: Description
- Functions: Description
"""

View file

@ -0,0 +1,7 @@
"""
Module specific exceptions for <this module>
Exceptions:
- List: Description
- Exceptions: Description
"""

View file

@ -0,0 +1,7 @@
"""
Database models for <this module>
Models:
- List: Description
- Models: Description
"""

View file

@ -0,0 +1,13 @@
"""
Router endpoints for <this module>
Endpoints:
- List: Description
- Endpoints: Description
"""
from fastapi import APIRouter
_router = APIRouter(
tags=[""],
)

View file

@ -0,0 +1,7 @@
"""
Pydantic models for <this module>
Models:
- List: Description
- Models: Description
"""

View file

@ -0,0 +1,11 @@
"""
Module specific business logic for <this module>
Classes:
- List: Description
- Classes: Description
Functions:
- List: Description
- Functions: Description
"""

View file

@ -0,0 +1,11 @@
"""
Non-business logic reusable functions and classes for <this module>
Classes:
- List: Description
- Classes: Description
Functions:
- List: Description
- Functions: Description
"""

21
src/api.py Normal file
View file

@ -0,0 +1,21 @@
"""
This module hooks the routers for the main endpoints into a single router for importing to the app.
"""
from fastapi import APIRouter
from src.auth.router import router as auth_router
from src.misp.router import router as misp_router
from src.control.router import router as control_router
api_router = APIRouter()
api_router.include_router(auth_router)
api_router.include_router(misp_router)
api_router.include_router(control_router)
@api_router.get("/healthcheck", include_in_schema=False)
def healthcheck():
"""Simple health check endpoint."""
return {"status": "ok"}

13
src/auth/config.py Normal file
View file

@ -0,0 +1,13 @@
"""
Configurations for auth module, import auth_settings
"""
from src.config import CustomBaseSettings
class AuthConfig(CustomBaseSettings):
OIDC_CONFIG: str = ""
OIDC_ISSUER: str = ""
OIDC_AUDIENCE: str = ""
CLIENT_ID: str = ""
auth_settings = AuthConfig()

7
src/auth/constants.py Normal file
View file

@ -0,0 +1,7 @@
"""
Constants and error codes for auth module
Constants:
- List: Description
- Consts: Description
"""

11
src/auth/dependencies.py Normal file
View file

@ -0,0 +1,11 @@
"""
Router dependencies for auth module
Classes:
- List: Description
- Classes: Description
Functions:
- List: Description
- Functions: Description
"""

7
src/auth/exceptions.py Normal file
View file

@ -0,0 +1,7 @@
"""
Module specific exceptions for auth module
Exceptions:
- List: Description
- Exceptions: Description
"""

7
src/auth/models.py Normal file
View file

@ -0,0 +1,7 @@
"""
Database models for auth module
Models:
- List: Description
- Models: Description
"""

11
src/auth/router.py Normal file
View file

@ -0,0 +1,11 @@
"""
Router endpoints for auth module
Contains oauth registration
Endpoints:
"""
from fastapi import APIRouter
router = APIRouter(
tags=["auth"],
)

7
src/auth/schemas.py Normal file
View file

@ -0,0 +1,7 @@
"""
Pydantic models for auth module
Models:
- List: Description
- Models: Description
"""

63
src/auth/service.py Normal file
View file

@ -0,0 +1,63 @@
"""
Module specific business logic for auth module
Exports:
- claims_dependency
- authed_dependency
"""
import json
from typing import Annotated
from authlib.jose import jwt
from urllib.request import urlopen
from fastapi import Depends, HTTPException
from fastapi.security import OpenIdConnect
from authlib.jose.rfc7517.jwk import JsonWebKey
from authlib.jose.rfc7517.key_set import KeySet
from authlib.oauth2.rfc7523.validator import JWTBearerToken
from src.auth.config import auth_settings
oidc = OpenIdConnect(openIdConnectUrl=auth_settings.OIDC_CONFIG)
oidc_dependency = Annotated[str, Depends(oidc)]
async def get_current_user(oidc_auth_string: oidc_dependency) -> JWTBearerToken:
config_url = urlopen(auth_settings.OIDC_CONFIG)
config = json.loads(config_url.read())
jwks_uri = config["jwks_uri"]
key_response = urlopen(jwks_uri)
jwk_keys: KeySet = JsonWebKey.import_key_set(json.loads(key_response.read()))
claims_options = {
"exp": {"essential": True},
"aud": {"essential": True, "value": "account"},
"iss": {"essential": True, "value": auth_settings.OIDC_ISSUER},
}
claims: JWTBearerToken = jwt.decode(
oidc_auth_string.replace("Bearer ", ""),
jwk_keys,
claims_options=claims_options,
claims_cls=JWTBearerToken,
)
claims.validate()
return claims
claims_dependency = Annotated[JWTBearerToken, Depends(get_current_user)]
async def is_authed_user(claims: claims_dependency) -> bool:
authed_users: list[str] = ["chris@sr2.uk"]
user_email = claims.get("email", None)
if not user_email or user_email not in authed_users:
raise HTTPException(status_code=403, detail="Not authenticated")
return claims.get("email") in authed_users
authed_dependency = Annotated[bool, Depends(is_authed_user)]

11
src/auth/utils.py Normal file
View file

@ -0,0 +1,11 @@
"""
Non-business logic reusable functions and classes for auth module
Classes:
- List: Description
- Classes: Description
Functions:
- List: Description
- Functions: Description
"""

53
src/config.py Normal file
View file

@ -0,0 +1,53 @@
"""
Global configurations: import settings, app_configs
Classes:
- CustomBaseSettings - Base class to be used by all modules for loading configs
"""
from typing import Any
from urllib import parse
from pydantic_settings import BaseSettings, SettingsConfigDict
from pydantic import SecretStr
from src.constants import Environment
class CustomBaseSettings(BaseSettings):
model_config = SettingsConfigDict(
env_file=".env", env_file_encoding="utf-8", extra="ignore"
)
class Config(CustomBaseSettings):
APP_VERSION: str = "0.1"
ENVIRONMENT: Environment = Environment.PRODUCTION
SECRET_KEY: SecretStr = ""
CORS_ORIGINS: list[str] = ["*"]
CORS_ORIGINS_REGEX: str | None = None
CORS_HEADERS: list[str] = ["*"]
DATABASE_NAME: str = "dns-exp"
DATABASE_PORT: str = "5432"
DATABASE_HOSTNAME: str = "localhost"
DATABASE_CREDENTIALS: SecretStr = ""
settings = Config()
DATABASE_NAME = settings.DATABASE_NAME
DATABASE_PORT = settings.DATABASE_PORT
DATABASE_HOSTNAME = settings.DATABASE_HOSTNAME
DATABASE_CREDENTIALS = settings.DATABASE_CREDENTIALS.get_secret_value()
# this will support special chars for credentials
_DATABASE_CREDENTIAL_USER, _DATABASE_CREDENTIAL_PASSWORD = str(DATABASE_CREDENTIALS).split(":")
_QUOTED_DATABASE_PASSWORD = parse.quote_plus(str(_DATABASE_CREDENTIAL_PASSWORD))
SQLALCHEMY_DATABASE_URI = SecretStr(f"postgresql+psycopg://{_DATABASE_CREDENTIAL_USER}:{_QUOTED_DATABASE_PASSWORD}@{DATABASE_HOSTNAME}:{DATABASE_PORT}/{DATABASE_NAME}")
app_configs: dict[str, Any] = {"title": "App API"}
if settings.ENVIRONMENT.is_deployed:
app_configs["root_path"] = f"/v{settings.APP_VERSION}"
if not settings.ENVIRONMENT.is_debug:
app_configs["openapi_url"] = None # hide docs

36
src/constants.py Normal file
View file

@ -0,0 +1,36 @@
"""
Global constants
Classes:
- Environment(StrEnum): LOCAL, TESTING, STAGING, PRODUCTION
"""
from enum import StrEnum, auto
class Environment(StrEnum):
"""
Enumeration of environments.
Attributes:
LOCAL (str): Application is running locally
TESTING (str): Application is running in testing mode
STAGING (str): Application is running in staging mode (ie not testing)
PRODUCTION (str): Application is running in production mode
"""
LOCAL = auto()
TESTING = auto()
STAGING = auto()
PRODUCTION = auto()
@property
def is_debug(self):
return self in (self.LOCAL, self.STAGING, self.TESTING)
@property
def is_testing(self):
return self == self.TESTING
@property
def is_deployed(self) -> bool:
return self in (self.STAGING, self.PRODUCTION)

7
src/control/config.py Normal file
View file

@ -0,0 +1,7 @@
"""
Configurations for <this module>
Configurations:
- List: Description
- Configs: Description
"""

13
src/control/constants.py Normal file
View file

@ -0,0 +1,13 @@
"""
Constants and error codes for <this module>
Constants:
- List: Description
- Consts: Description
"""
from enum import StrEnum, auto
class TimerState(StrEnum):
STARTING = auto()
STOPPING = auto()

View file

@ -0,0 +1,11 @@
"""
Router dependencies for <this module>
Classes:
- List: Description
- Classes: Description
Functions:
- List: Description
- Functions: Description
"""

View file

@ -0,0 +1,7 @@
"""
Module specific exceptions for <this module>
Exceptions:
- List: Description
- Exceptions: Description
"""

7
src/control/models.py Normal file
View file

@ -0,0 +1,7 @@
"""
Database models for <this module>
Models:
- List: Description
- Models: Description
"""

34
src/control/router.py Normal file
View file

@ -0,0 +1,34 @@
"""
Router endpoints for the control module
Endpoints:
- List: Description
- Endpoints: Description
"""
from fastapi import APIRouter, Request
from starlette import status
from src.auth.service import authed_dependency
from src.control.schemas import ControlTimerPutResponse
router = APIRouter(
tags=["control"],
prefix="/control",
)
@router.put("/start_timer", status_code=status.HTTP_202_ACCEPTED, response_model=ControlTimerPutResponse)
async def start_timer(request: Request):
misp_handler = request.app.misp_handler
await misp_handler.start_timer()
return {"state": "starting"}
@router.put("/stop_timer", status_code=status.HTTP_202_ACCEPTED, response_model=ControlTimerPutResponse)
async def stop_timer(request: Request):
misp_handler = request.app.misp_handler
misp_handler.stop_timer()
return {"state": "stopping"}

14
src/control/schemas.py Normal file
View file

@ -0,0 +1,14 @@
"""
Pydantic models for <this module>
Models:
- List: Description
- Models: Description
"""
from src.schemas import CustomBaseModel
from src.control.constants import TimerState
class ControlTimerPutResponse(CustomBaseModel):
state: TimerState

11
src/control/service.py Normal file
View file

@ -0,0 +1,11 @@
"""
Module specific business logic for <this module>
Classes:
- List: Description
- Classes: Description
Functions:
- List: Description
- Functions: Description
"""

11
src/control/utils.py Normal file
View file

@ -0,0 +1,11 @@
"""
Non-business logic reusable functions and classes for <this module>
Classes:
- List: Description
- Classes: Description
Functions:
- List: Description
- Functions: Description
"""

31
src/database.py Normal file
View file

@ -0,0 +1,31 @@
"""
Database connections and init
Exports:
- db_dependency
- Base (sqlalchemy base model)
- get_db
"""
from typing import Annotated
from sqlalchemy import create_engine
from sqlalchemy.orm import declarative_base, sessionmaker, Session
from fastapi import Depends
from src.config import SQLALCHEMY_DATABASE_URI
engine = create_engine(SQLALCHEMY_DATABASE_URI.get_secret_value())
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
def get_db():
with SessionLocal.begin() as db:
try:
yield db
finally:
db.rollback() # Anything not explicitly commited is rolled back
db.close()
db_dependency = Annotated[Session, Depends(get_db)]
Base = declarative_base()

3
src/exceptions.py Normal file
View file

@ -0,0 +1,3 @@
"""
Global exceptions
"""

67
src/main.py Normal file
View file

@ -0,0 +1,67 @@
"""
Application root file: Inits the FastAPI application
"""
from contextlib import asynccontextmanager
from typing import AsyncGenerator
from prometheus_client import make_asgi_app
from fastapi import FastAPI
from starlette.middleware.sessions import SessionMiddleware
from starlette.middleware.cors import CORSMiddleware
from src.config import settings
from src.api import api_router
from src.prometheus import prometheus
from src.auth.config import auth_settings
from src.misp.service import MISPHandler
# TODO: Create Pydantic request/response schemas
@asynccontextmanager
async def lifespan(_application: FastAPI) -> AsyncGenerator:
# Startup
yield
# Shutdown
_application.misp_handler.stop_timer()
if settings.ENVIRONMENT.is_deployed:
# Do this only on prod
pass
app = FastAPI(
swagger_ui_init_oauth={
"clientId": auth_settings.CLIENT_ID,
"usePkceWithAuthorizationCodeGrant": True,
"scopes": "openid profile email",
}
)
metrics_app = make_asgi_app()
app.mount("/metrics", metrics_app)
prometheus.APP_STATE.state("starting")
# Type inspection disabled for middleware injection.
# Known bug in FastAPI type checking: https://github.com/astral-sh/ty/issues/1635
# noinspection PyTypeChecker
app.add_middleware(SessionMiddleware, secret_key=settings.SECRET_KEY.get_secret_value())
# noinspection PyTypeChecker
app.add_middleware(
CORSMiddleware,
allow_origins=settings.CORS_ORIGINS,
allow_origin_regex=settings.CORS_ORIGINS_REGEX,
allow_credentials=True,
allow_methods=("GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS"),
allow_headers=settings.CORS_HEADERS,
)
app.include_router(api_router)
app.misp_handler = MISPHandler()
prometheus.APP_STATE.state("running")

14
src/misp/config.py Normal file
View file

@ -0,0 +1,14 @@
"""
Configurations for the MISP module
"""
from src.config import CustomBaseSettings
class MISPConfig(CustomBaseSettings):
MISP_KEY_FILE: str = ""
MISP_OUTPUT_FILE: str = ""
ALLOWED_TLP: list[str] = ["tlp:clear", "tlp:white", "tlp:green"]
IGNORED_TLP: list[str] = ["tlp:red", "tlp:amber+strict", "tlp:amber"]
UNBOUND_CERT_DIR: str = ""
settings = MISPConfig()

7
src/misp/constants.py Normal file
View file

@ -0,0 +1,7 @@
"""
Constants and error codes for <this module>
Constants:
- List: Description
- Consts: Description
"""

11
src/misp/dependencies.py Normal file
View file

@ -0,0 +1,11 @@
"""
Router dependencies for <this module>
Classes:
- List: Description
- Classes: Description
Functions:
- List: Description
- Functions: Description
"""

7
src/misp/exceptions.py Normal file
View file

@ -0,0 +1,7 @@
"""
Module specific exceptions for <this module>
Exceptions:
- List: Description
- Exceptions: Description
"""

18
src/misp/models.py Normal file
View file

@ -0,0 +1,18 @@
"""
Database models for the misp module
"""
from sqlalchemy import Column, Integer, String, Boolean
from sqlalchemy.dialects.postgresql import ARRAY
from sqlalchemy.ext.mutable import MutableList
from src.database import Base
class Domain(Base):
__tablename__ = "domain"
id = Column(Integer, primary_key=True)
domain = Column(String, index=True)
events = Column(MutableList.as_mutable(ARRAY(Integer)), default=[])
always_allowed = Column(Boolean, default=False)
ignored_events = Column(MutableList.as_mutable(ARRAY(Integer)), default=[])

123
src/misp/router.py Normal file
View file

@ -0,0 +1,123 @@
"""
Router endpoints for <this module>
Endpoints:
- List: Description
- Endpoints: Description
"""
from fastapi import APIRouter, HTTPException, Request, BackgroundTasks
from sqlalchemy.sql import and_
from src.auth.service import authed_dependency
from src.database import db_dependency
from src.misp.models import Domain
from src.misp.schemas import MispUpdatePutRequest, MispUpdatePutResponse
router = APIRouter(
tags=["misp"],
prefix="/misp",
)
@router.put("/manual_update", response_model=MispUpdatePutResponse)
async def manual_misp_update(request: Request, update_request: MispUpdatePutRequest, background_tasks: BackgroundTasks):
published_time = update_request.published_timestamp
misp_handler = request.app.misp_handler
background_tasks.add_task(misp_handler.full_update_once, p_time=published_time)
if not published_time:
published_time = "default"
return {"time": published_time, "state": "Starting"}
@router.get("/domain/blocked/{domain}")
async def domain_blocked(domain: str, db: db_dependency):
same_elements = and_(
Domain.events.contains(Domain.ignored_events),
Domain.events.contained_by(Domain.ignored_events)
)
domain_model = (db.query(Domain).filter(Domain.domain == domain,
~Domain.always_allowed,
~same_elements
).first())
return {"is_blocked": bool(domain_model)}
@router.get("/domain/search")
async def domain_search(domain: str, db: db_dependency):
domain = domain.replace("*", "%")
same_elements = and_(
Domain.events.contains(Domain.ignored_events),
Domain.events.contained_by(Domain.ignored_events)
)
query = db.query(Domain, same_elements.label("is_ignored)")).filter(Domain.domain.ilike(domain))
results = query.all()
def domain_status(item_details):
if item_details[0].always_allowed:
return "allowed"
elif item_details[1]:
return "ignored"
else:
return "blocked"
return {item[0].domain: domain_status(item) for item in results}
@router.patch("/domain/always_allowed/{domain}")
async def always_allowed(db: db_dependency, domain: str, allow: bool):
domain_model = db.query(Domain).filter(Domain.domain == domain).first()
if domain_model:
domain_model.always_allowed = allow
else:
domain_model = Domain(domain=domain, events=[], always_allowed=allow)
db.add(domain_model)
db.commit()
@router.patch("/domain/events/{domain}/ignore")
async def event_ignore(domain: str, db: db_dependency, event: int):
domain_model = db.query(Domain).filter(Domain.domain == domain).first()
if not domain_model:
raise HTTPException(status_code=404, detail="Domain Not Found")
if event not in domain_model.events:
raise HTTPException(status_code=404, detail="Event Not Found")
ignored_events = domain_model.ignored_events or []
if event in ignored_events:
return {"status": "Event Ignored"}
domain_model.ignored_events = ignored_events + [event]
db.commit()
return {"status": "Event Ignored"}
@router.patch("/domain/events/{domain}/reinstate")
async def event_reinstate(domain: str, db: db_dependency, event: int):
domain_model = db.query(Domain).filter(Domain.domain == domain).first()
if not domain_model:
raise HTTPException(status_code=404, detail="Domain Not Found")
ignored_events = domain_model.ignored_events
if not ignored_events or event not in ignored_events:
return {"status": "Event Ignored"}
ignored_events.remove(event)
domain_model.ignored_events = ignored_events
db.add(domain_model)
db.commit()
return {"status": "Event Un-ignored"}

20
src/misp/schemas.py Normal file
View file

@ -0,0 +1,20 @@
"""
Pydantic models for <this module>
Models:
- List: Description
- Models: Description
"""
from typing import Optional
from pydantic import Field
from src.schemas import CustomBaseModel
class MispUpdatePutRequest(CustomBaseModel):
published_timestamp: Optional[str] = Field(default=None, description="Timestamp for how far back to check for published timestamps")
class MispUpdatePutResponse(CustomBaseModel):
time: str
state: str

268
src/misp/service.py Normal file
View file

@ -0,0 +1,268 @@
"""
Module specific business logic for the MISP module
Classes:
- List: Description
- Classes: Description
Functions:
- List: Description
- Functions: Description
"""
import os
import json
import threading
from typing import Any, Optional
from sqlalchemy.sql import and_
from pymisp import PyMISP
from unbound_console import RemoteControl
from src.config import settings as global_settings
from src.misp.config import settings as misp_settings
from src.database import get_db
from src.misp.models import Domain
from src.prometheus import prometheus
class MISPHandler:
def __init__(self) -> None:
self.CONFIG_HEADER: str = """server:
interface: 0.0.0.0@853
interface: ::0@853
interface: 0.0.0.0@443
interface: ::0@443
interface: 0.0.0.0@53
interface: ::0@53
tls-service-key: "/etc/letsencrypt/live/dns.sr2.uk/privkey.pem"
tls-service-pem: "/etc/letsencrypt/live/dns.sr2.uk/fullchain.pem"
tls-port: 853
https-port: 443
access-control: 0.0.0.0/0 allow
access-control: ::/0 allow
do-ip4: yes
do-ip6: yes
do-udp: yes
do-tcp: yes
prefetch: yes
"""
self.scheduled_published_timestamp = "7d"
self.domains: dict[str, set[int]] = {}
self.timer: Optional[threading.Thread] = None
self._stop_event: threading.Event = threading.Event()
self.is_fetching: bool = False
self.dev_dump: bool = True
self.timer_time: int = 30
self.pymisp = PyMISP(
url="https://misp.civicert.org",
key=open(misp_settings.MISP_KEY_FILE, "r").read().strip(),
)
prometheus.TIMER_STATE.state("stopped")
prometheus.UNBOUND_RELOAD_RESPONSE.state("not reloaded")
self.update_metrics()
async def start_timer(self) -> None:
if self.timer is not None and self.timer.is_alive():
print("Timer already started")
return
self._stop_event.clear()
self.timer = threading.Thread(target=self.run_with_timer, daemon=True)
if self.timer:
self.timer.start()
def run_with_timer(self) -> None:
prometheus.TIMER_STATE.state("running")
try:
while not self._stop_event.wait(self.timer_time):
self.full_update_once()
finally:
prometheus.TIMER_STATE.state("stopped")
self.timer = None
self._stop_event.clear()
def stop_timer(self) -> None:
if self.timer is None:
return
prometheus.TIMER_STATE.state("stopping")
self._stop_event.set()
self.timer.join()
self.timer = None
@staticmethod
def is_event_tlp_ok(event: dict[str, Any]) -> bool:
for tag in event.get("Tag", []):
if tag["name"] in misp_settings.ALLOWED_TLP:
return True
return False
def is_tlp_ok(self, attribute: dict[str, Any]) -> bool:
for tag in attribute.get("Tag", []):
if tag["name"] in misp_settings.IGNORED_TLP:
return False
if tag["name"] in misp_settings.ALLOWED_TLP:
return True
return self.is_event_tlp_ok(attribute["Event"])
def add_domains_to_db(self) -> None:
prometheus.MISP_STATE.state("updating")
db = next(get_db())
for domain, events in self.domains.items():
domain_model = (db.query(Domain).filter(Domain.domain == domain).first())
if domain_model:
existing_events = set(domain_model.events)
existing_events.update(events)
domain_model.events = list(existing_events)
else:
domain_model = Domain(domain=domain, events=list(events), always_allowed=False, ignored_events=[])
db.add(domain_model)
db.commit()
prometheus.MISP_STATE.state("idle")
def get_domains(self, type_: str, page: int, p_time: Optional[str]) -> tuple[dict[str, set[int]], int]:
if p_time is None:
p_time = self.scheduled_published_timestamp
# pymisp.search with these parameters returns a dict. Other params may return a different structure so inspection flags a type error.
# noinspection PyTypeChecker
data: dict[str, Any] = self.pymisp.search(
controller="attributes",
type_attribute=type_,
page=page,
limit=1000,
publish_timestamp=p_time,
published=True,
to_ids=True,
deleted="0",
include_event_tags=True
)
if self.dev_dump:
with open("data_dump.json", "w", encoding="utf-8") as dump_file:
json.dump(data, dump_file, indent=4)
self.dev_dump = False
domains = {}
for attribute in data["Attribute"]:
if not self.is_tlp_ok(attribute):
continue
domain = attribute["value"]
event_id = int(attribute["Event"]["id"])
if domain not in domains:
domains[domain] = {event_id}
else:
domains[domain].add(event_id)
return domains, len(data["Attribute"])
def produce_domain_set(self, p_time: Optional[str]) -> None:
self.is_fetching = True
prometheus.MISP_STATE.state("fetching")
for type_ in ["domain", "hostname"]:
page = 1
while True:
page_domains, page_size = self.get_domains(type_, page, p_time)
if not page_size:
break
for key, value in page_domains.items():
if key in self.domains:
self.domains[key].update(value)
else:
self.domains[key] = value
print(type_, page, len(page_domains), len(self.domains.keys()))
page += 1
if global_settings.ENVIRONMENT.is_debug:
break
self.is_fetching = False
prometheus.MISP_STATE.state("idle")
def save_conf(self) -> None:
db = next(get_db())
same_elements = and_(
Domain.events.contains(Domain.ignored_events),
Domain.events.contained_by(Domain.ignored_events)
)
results = (db.query(Domain)
.filter(
~Domain.always_allowed,
~same_elements
)
.all()
)
blocked_domains = [item.domain for item in results]
with open(misp_settings.MISP_OUTPUT_FILE, "w") as conf:
conf.write(self.CONFIG_HEADER)
for domain in blocked_domains:
conf.write(f' local-zone: "{domain}" always_nxdomain\n')
conf.write(' local-zone: "stats.invalid." static\n')
conf.write(f' local-data: "stats.invalid. TXT \'{len(self.domains.keys())} domains in blocklist"\n')
@staticmethod
def send_unbound_command(command) -> str:
rc = RemoteControl(
host=misp_settings.UNBOUND_HOST,
port=misp_settings.UNBOUND_PORT,
server_cert=os.path.join(misp_settings.UNBOUND_CERT_DIR, "unbound_server.pem"),
client_cert=os.path.join(misp_settings.UNBOUND_CERT_DIR, "unbound_control.pem"),
client_key=os.path.join(misp_settings.UNBOUND_CERT_DIR, "unbound_control.key"),
)
return rc.send_command(cmd=command)
def reload_service(self) -> None:
prometheus.MISP_STATE.state("reloading service")
response = self.send_unbound_command("reload")
if response:
if response == "0":
prometheus.UNBOUND_RELOAD_RESPONSE.state("success")
elif response == "1":
prometheus.UNBOUND_RELOAD_RESPONSE.state("error")
else:
prometheus.UNBOUND_RELOAD_RESPONSE.state("no response")
prometheus.MISP_STATE.state("idle")
@staticmethod
def update_metrics() -> None:
db = next(get_db())
same_elements = and_(
Domain.events.contains(Domain.ignored_events),
Domain.events.contained_by(Domain.ignored_events)
)
blocked_count = (db.query(Domain)
.filter(
~Domain.always_allowed,
~same_elements
)
.count()
)
prometheus.blocked_domain_count.set(blocked_count)
@prometheus.last_update_length.time()
def full_update_once(self, p_time: Optional[str] = None) -> None:
if self.is_fetching:
return
self.produce_domain_set(p_time)
self.add_domains_to_db()
self.save_conf()
if not global_settings.ENVIRONMENT.is_debug:
self.reload_service()
prometheus.last_update_complete_time.set_to_current_time()
self.update_metrics()
print("Done")

11
src/misp/utils.py Normal file
View file

@ -0,0 +1,11 @@
"""
Non-business logic reusable functions and classes for <this module>
Classes:
- List: Description
- Classes: Description
Functions:
- List: Description
- Functions: Description
"""

4
src/models.py Normal file
View file

@ -0,0 +1,4 @@
"""
Global database models
"""

View file

View file

@ -0,0 +1,51 @@
from prometheus_client import Enum, Gauge
APP_STATE = Enum(
"app_state",
"Current state of the application",
states=["starting", "running", "error"],
namespace="app"
)
MISP_STATE = Enum(
"misp_state",
"Current state of the ",
states=["idle", "fetching", "updating", "reloading service", "error"],
namespace="misp"
)
TIMER_STATE = Enum(
"timer_state",
"Current state of the timer",
states=["running", "stopping", "stopped"],
namespace="misp",
subsystem="timer"
)
UNBOUND_RELOAD_RESPONSE = Enum(
"unbound_reload_response",
"Response from last unbound reload command",
states=["success", "error", "no response", "not reloaded"],
namespace="misp",
subsystem="unbound"
)
blocked_domain_count = Gauge(
"blocked_domain_count",
"Number of blocked domains",
namespace="misp",
subsystem="domains"
)
last_update_length = Gauge(
"last_update_time",
"How long the last update took",
namespace="misp",
)
last_update_complete_time = Gauge(
"last_update_complete_time",
"The time at which the last update was completed",
namespace="misp",
)

5
src/schemas.py Normal file
View file

@ -0,0 +1,5 @@
from pydantic import BaseModel
class CustomBaseModel(BaseModel):
pass

98
src/utils.py Normal file
View file

@ -0,0 +1,98 @@
import asyncio
import logging
import inspect
from functools import wraps
from traceback import format_exception
from typing import Coroutine, Callable, Any
from starlette.concurrency import run_in_threadpool
NoArgsNoReturnFuncT = Callable[[], None]
NoArgsNoReturnAsyncFuncT = Callable[[], Coroutine[Any, Any, None]]
ExcArgNoReturnFuncT = Callable[[Exception], None]
ExcArgNoReturnAsyncFuncT = Callable[[Exception], Coroutine[Any, Any, None]]
NoArgsNoReturnAnyFuncT = NoArgsNoReturnFuncT | NoArgsNoReturnAsyncFuncT
ExcArgNoReturnAnyFuncT = ExcArgNoReturnFuncT | ExcArgNoReturnAsyncFuncT
NoArgsNoReturnDecorator = Callable[[NoArgsNoReturnAnyFuncT], NoArgsNoReturnAsyncFuncT]
async def _handle_repeat_func(func: NoArgsNoReturnAnyFuncT) -> None:
if inspect.iscoroutinefunction(func):
await func()
else:
await run_in_threadpool(func)
async def _handle_repeat_exc(
exc: Exception, on_exception: ExcArgNoReturnAnyFuncT | None
) -> None:
if on_exception:
if inspect.iscoroutinefunction(on_exception):
await on_exception(exc)
else:
await run_in_threadpool(on_exception, exc)
def repeat_every(
*,
seconds: float,
wait_first: float | None = None,
max_repetitions: int | None = None,
on_complete: NoArgsNoReturnAnyFuncT | None = None,
on_exception: ExcArgNoReturnAnyFuncT | None = None,
) -> NoArgsNoReturnDecorator:
"""
This function returns a decorator that modifies a function so it is periodically re-executed after its first call.
The function it decorates should accept no arguments and return nothing. If necessary, this can be accomplished
by using `functools.partial` or otherwise wrapping the target function prior to decoration.
Parameters
----------
seconds: float
The number of seconds to wait between repeated calls
wait_first: float (default None)
If not None, the function will wait for the given duration before the first call
max_repetitions: Optional[int] (default None)
The maximum number of times to call the repeated function. If `None`, the function is repeated forever.
on_complete: Optional[Callable[[], None]] (default None)
A function to call after the final repetition of the decorated function.
on_exception: Optional[Callable[[Exception], None]] (default None)
A function to call when an exception is raised by the decorated function.
"""
def decorator(func: NoArgsNoReturnAnyFuncT) -> NoArgsNoReturnAsyncFuncT:
"""
Converts the decorated function into a repeated, periodically-called version of itself.
"""
@wraps(func)
async def wrapped() -> None:
async def loop() -> None:
if wait_first is not None:
await asyncio.sleep(wait_first)
repetitions = 0
while max_repetitions is None or repetitions < max_repetitions:
try:
await _handle_repeat_func(func)
except Exception as exc:
formatted_exception = "".join(
format_exception(type(exc), exc, exc.__traceback__)
)
logging.error(formatted_exception)
await _handle_repeat_exc(exc, on_exception)
repetitions += 1
await asyncio.sleep(seconds)
if on_complete:
await _handle_repeat_func(on_complete)
asyncio.ensure_future(loop())
return wrapped
return decorator