From 9aa390433b79b41343db9739be4e58b6b559c785 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Kolankiewicz?= Date: Tue, 24 Sep 2024 15:13:35 +0200 Subject: [PATCH 01/11] Issue #401: bump backend dependencies --- requirements.txt | 55 ++++++++++++++++++++++++------------------------ 1 file changed, 28 insertions(+), 27 deletions(-) diff --git a/requirements.txt b/requirements.txt index 46389868..f3181072 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,28 +1,29 @@ -anyio==3.4.0 -asgiref==3.4.1 -cachetools==4.2.4 -certifi==2021.10.8 -charset-normalizer==2.0.9 -click==8.0.3 -Deprecated==1.2.13 -fastapi==0.70.0 -h11==0.12.0 -idna==3.3 -pydantic==1.10.14 -pyzmq==24.0.1 -redis==4.0.2 -requests==2.26.0 -sniffio==1.2.0 -starlette==0.16.0 -typing-extensions==4.9.0 -urllib3==1.26.7 -uvicorn==0.15.0 -wrapt==1.13.3 -yara-python==4.1.3 -yaramod==3.12.1 -PyJWT[crypto]==2.3.0 -rq==1.11.1 -typed-config==1.3.2 -sqlmodel==0.0.11 +annotated-types==0.7.0 +anyio==4.6.0 +cachetools==5.5.0 +certifi==2024.8.30 +cffi==1.17.1 +charset-normalizer==3.3.2 +click==8.1.7 +cryptography==43.0.1 +fastapi==0.115.0 +h11==0.14.0 +idna==3.10 psycopg2==2.9.9 -alembic==1.11.1 +pycparser==2.22 +pydantic==1.10.18 +pydantic_core==2.23.4 +PyJWT==2.9.0 +pyzmq==26.2.0 +redis==5.0.8 +requests==2.32.3 +rq==1.16.2 +sniffio==1.3.1 +sqlmodel==0.0.22 +starlette==0.38.6 +typed-config==2.0.3 +typing_extensions==4.12.2 +urllib3==2.2.3 +uvicorn==0.30.6 +yara-python==4.5.1 +yaramod==3.23.0 \ No newline at end of file From 4db91da787edaac8887ccad8e9726f4d6d0a32cc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Kolankiewicz?= Date: Tue, 24 Sep 2024 15:22:41 +0200 Subject: [PATCH 02/11] fix: missing packages added --- requirements.txt | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index f3181072..d5d4ad79 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,11 +1,14 @@ +alembic==1.11.1 annotated-types==0.7.0 anyio==4.6.0 +asgiref==3.4.1 cachetools==5.5.0 certifi==2024.8.30 cffi==1.17.1 charset-normalizer==3.3.2 click==8.1.7 cryptography==43.0.1 +Deprecated==1.2.13 fastapi==0.115.0 h11==0.14.0 idna==3.10 @@ -13,7 +16,7 @@ psycopg2==2.9.9 pycparser==2.22 pydantic==1.10.18 pydantic_core==2.23.4 -PyJWT==2.9.0 +PyJWT[crypto]==2.9.0 pyzmq==26.2.0 redis==5.0.8 requests==2.32.3 @@ -25,5 +28,6 @@ typed-config==2.0.3 typing_extensions==4.12.2 urllib3==2.2.3 uvicorn==0.30.6 +wrapt==1.16.0 yara-python==4.5.1 yaramod==3.23.0 \ No newline at end of file From 09d7fb5a7bf826cace1a9fdc9e2d79e4b255d84f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Kolankiewicz?= Date: Mon, 30 Sep 2024 08:30:25 +0200 Subject: [PATCH 03/11] lint: mypy ignores site-packages directory --- setup.cfg | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.cfg b/setup.cfg index f56b362d..b4de06f9 100644 --- a/setup.cfg +++ b/setup.cfg @@ -4,6 +4,7 @@ exclude = mqueryfront/ [mypy] python_version = 3.6 +no_site_packages = True [mypy-yaramod.*] ignore_missing_imports = True From bdd5b1ee585a7e748285e841f59668809f39bf19 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Kolankiewicz?= Date: Mon, 30 Sep 2024 09:32:33 +0200 Subject: [PATCH 04/11] lint: bump python version --- setup.cfg | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/setup.cfg b/setup.cfg index b4de06f9..2c4a37d2 100644 --- a/setup.cfg +++ b/setup.cfg @@ -3,8 +3,7 @@ ignore = E501,W503,E203 exclude = mqueryfront/ [mypy] -python_version = 3.6 -no_site_packages = True +python_version = 3.10 [mypy-yaramod.*] ignore_missing_imports = True From fa76c06fa85435ec83469b07e0aae92fc6f9ec05 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Kolankiewicz?= Date: Mon, 30 Sep 2024 09:50:45 +0200 Subject: [PATCH 05/11] lint: bump mypy version and add ignore types --- .github/workflows/test_code.yml | 4 ++-- CONTRIBUTING.md | 2 +- src/lib/ursadb.py | 4 ++-- src/lib/yaraparse.py | 20 +++++++++---------- src/models/agentgroup.py | 2 +- src/models/configentry.py | 2 +- src/models/job.py | 2 +- src/models/jobagent.py | 2 +- src/models/match.py | 2 +- src/tests/test_ursadb.py | 3 +-- src/tests/test_yaraparse.py | 2 +- .../generate_yaraparse_result_files.py | 4 ++-- src/tests/yararules/test_corpus_yara_rule.py | 4 ++-- src/utils/compactall.py | 2 +- src/utils/index.py | 2 +- src/utils/s3index.py | 2 +- 16 files changed, 29 insertions(+), 30 deletions(-) diff --git a/.github/workflows/test_code.yml b/.github/workflows/test_code.yml index ff1aae71..a31074c5 100644 --- a/.github/workflows/test_code.yml +++ b/.github/workflows/test_code.yml @@ -24,8 +24,8 @@ jobs: uses: actions/setup-python@v1 with: python-version: '3.10' - - name: install mypy==0.790 - run: pip3 install mypy==0.790 + - name: install mypy==0.981 + run: pip3 install mypy==0.981 - name: install requirements run: pip3 install -r requirements.txt - name: run mypy diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 0a82fdb3..eccec677 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -50,7 +50,7 @@ $ prettier --write src/mqueryfront/ - Verify that there are no type errors with [mypy](http://mypy-lang.org/): ```bash -$ pip install mypy==0.790 +$ pip install mypy==0.981 $ mypy src ``` diff --git a/src/lib/ursadb.py b/src/lib/ursadb.py index 7e97125f..786ce8b7 100644 --- a/src/lib/ursadb.py +++ b/src/lib/ursadb.py @@ -48,12 +48,12 @@ def __execute(self, command: str, recv_timeout: int = 2000) -> Json: socket.send_string(command) return json.loads(socket.recv_string()) finally: - socket.close() + socket.close() # type: ignore def query( self, query: str, - taints: List[str] = None, + taints: List[str] = None, # type: ignore dataset: Optional[str] = None, ) -> Json: command = "select " diff --git a/src/lib/yaraparse.py b/src/lib/yaraparse.py index 44be487d..b6377425 100644 --- a/src/lib/yaraparse.py +++ b/src/lib/yaraparse.py @@ -377,7 +377,7 @@ def ursify_xor_string(string: PlainString) -> UrsaExpression: def ursify_string(string: String) -> Optional[UrsaExpression]: if string.is_xor: - return ursify_xor_string(string) + return ursify_xor_string(string) # type: ignore elif string.is_plain: return ursify_plain_string( string.pure_text, @@ -389,7 +389,7 @@ def ursify_string(string: String) -> Optional[UrsaExpression]: value_safe = string.pure_text.decode() return ursify_hex(value_safe) elif string.is_regexp: - return ursify_regex_string(string) + return ursify_regex_string(string) # type: ignore return None @@ -429,7 +429,7 @@ def pare_expr( def str_expr( self, condition: StringExpression ) -> Optional[UrsaExpression]: - return ursify_string(self.strings[condition.id]) + return ursify_string(self.strings[condition.id]) # type: ignore def expand_string_wildcard( self, condition: StringWildcardExpression @@ -440,7 +440,7 @@ def expand_string_wildcard( v for k, v in self.strings.items() if re.match(condition_regex, k) ] - ursa_strings = [ursify_string(x) for x in filtered_strings] + ursa_strings = [ursify_string(x) for x in filtered_strings] # type: ignore return [s for s in ursa_strings if s is not None] def expand_set_expression( @@ -463,7 +463,7 @@ def of_expr(self, condition: OfExpression) -> Optional[UrsaExpression]: if type(children) is SetExpression: all_elements = self.expand_set_expression(children) elif type(children) is ThemExpression: - all_elements = [ursify_string(k) for k in self.strings.values()] + all_elements = [ursify_string(k) for k in self.strings.values()] # type: ignore else: raise YaraParseError(f"Unsupported of_expr type: {type(children)}") @@ -524,7 +524,7 @@ def str_count_expr( self, condition: StringCountExpression ) -> Optional[UrsaExpression]: fixed_id = "$" + condition.id[1:] - return ursify_string(self.strings[fixed_id]) + return ursify_string(self.strings[fixed_id]) # type: ignore def int_lit_expr( self, condition: IntLiteralExpression @@ -535,15 +535,15 @@ def int_lit_expr( def str_at_expr( self, condition: StringAtExpression ) -> Optional[UrsaExpression]: - return ursify_string(self.strings[condition.id]) - + return ursify_string(self.strings[condition.id]) # type: ignore + def id_expr(self, condition: IdExpression) -> Optional[UrsaExpression]: return self.rules[condition.symbol.name].parse() def str_in_expr( self, condition: StringInRangeExpression ) -> Optional[UrsaExpression]: - return ursify_string(self.strings[condition.id]) + return ursify_string(self.strings[condition.id]) # type: ignore CONDITION_HANDLERS = { AndExpression: and_expr, @@ -565,7 +565,7 @@ def str_in_expr( def traverse(self, condition) -> Optional[UrsaExpression]: if type(condition) in self.CONDITION_HANDLERS: - return self.CONDITION_HANDLERS[type(condition)](self, condition) + return self.CONDITION_HANDLERS[type(condition)](self, condition) # type: ignore else: print(f"unsupported expression: {type(condition)}") return None diff --git a/src/models/agentgroup.py b/src/models/agentgroup.py index 70fa06a0..2325ea25 100644 --- a/src/models/agentgroup.py +++ b/src/models/agentgroup.py @@ -10,7 +10,7 @@ class AgentGroupBase(SQLModel): active_plugins: List[str] = Field(sa_column=Column(ARRAY(String))) -class AgentGroup(AgentGroupBase, table=True): +class AgentGroup(AgentGroupBase, table=True): # type: ignore """Agent group is a group of processes working on a single file group, with a shared storage, and a single backing ursadb. """ diff --git a/src/models/configentry.py b/src/models/configentry.py index 8e16256b..b9cb980b 100644 --- a/src/models/configentry.py +++ b/src/models/configentry.py @@ -2,7 +2,7 @@ from typing import Union -class ConfigEntry(SQLModel, table=True): +class ConfigEntry(SQLModel, table=True): # type: ignore id: Union[int, None] = Field(default=None, primary_key=True) plugin: str key: str diff --git a/src/models/job.py b/src/models/job.py index 17db42ad..2bc8f25d 100644 --- a/src/models/job.py +++ b/src/models/job.py @@ -30,7 +30,7 @@ class JobBase(SQLModel): agents_left: int -class Job(JobBase, table=True): +class Job(JobBase, table=True): # type: ignore """Job object in the database. Internal ID is an implementation detail.""" internal_id: Union[int, None] = Field(default=None, primary_key=True) diff --git a/src/models/jobagent.py b/src/models/jobagent.py index f40313de..e8a17564 100644 --- a/src/models/jobagent.py +++ b/src/models/jobagent.py @@ -6,7 +6,7 @@ from ..models.agentgroup import AgentGroup -class JobAgent(SQLModel, table=True): +class JobAgent(SQLModel, table=True): # type: ignore """Information about job relating to a specific agent group.""" id: Union[int, None] = Field(default=None, primary_key=True) diff --git a/src/models/match.py b/src/models/match.py index 23980822..e0fcc6cb 100644 --- a/src/models/match.py +++ b/src/models/match.py @@ -4,7 +4,7 @@ from ..models.job import Job -class Match(SQLModel, table=True): +class Match(SQLModel, table=True): # type: ignore """Represents a file matched to a job, along with a related metadata.""" id: Union[int, None] = Field(default=None, primary_key=True) diff --git a/src/tests/test_ursadb.py b/src/tests/test_ursadb.py index e628ec80..8de8a999 100644 --- a/src/tests/test_ursadb.py +++ b/src/tests/test_ursadb.py @@ -8,8 +8,7 @@ from typing import Dict, Any sys.path = [".."] + sys.path -from lib.ursadb import UrsaDb # noqa - +from lib.ursadb import UrsaDb # type: ignore class UrsadbTestContext: def __init__(self, socket: zmq.Socket, ursadb: UrsaDb): diff --git a/src/tests/test_yaraparse.py b/src/tests/test_yaraparse.py index c9bda913..3c9af2b8 100644 --- a/src/tests/test_yaraparse.py +++ b/src/tests/test_yaraparse.py @@ -1,6 +1,6 @@ """Unit tests for yaraparse.""" -from lib.yaraparse import ursify_hex, ursify_plain_string, parse_yara +from lib.yaraparse import ursify_hex, ursify_plain_string, parse_yara # type: ignore import yaramod diff --git a/src/tests/yararules/generate_yaraparse_result_files.py b/src/tests/yararules/generate_yaraparse_result_files.py index fa69f11f..54275526 100644 --- a/src/tests/yararules/generate_yaraparse_result_files.py +++ b/src/tests/yararules/generate_yaraparse_result_files.py @@ -1,7 +1,7 @@ import argparse import os -from lib.yaraparse import parse_yara -from lib.yaraparse import combine_rules +from lib.yaraparse import parse_yara # type: ignore +from lib.yaraparse import combine_rules # type: ignore current_path = os.path.abspath(os.path.dirname(__file__)) testdir = current_path + "/testdata/" diff --git a/src/tests/yararules/test_corpus_yara_rule.py b/src/tests/yararules/test_corpus_yara_rule.py index f34049f1..7e513e7a 100644 --- a/src/tests/yararules/test_corpus_yara_rule.py +++ b/src/tests/yararules/test_corpus_yara_rule.py @@ -2,8 +2,8 @@ import unittest from pathlib import Path -from lib.yaraparse import combine_rules -from lib.yaraparse import parse_yara +from lib.yaraparse import combine_rules # type: ignore +from lib.yaraparse import parse_yara # type: ignore testdir = Path(__file__).parent / "testdata" diff --git a/src/utils/compactall.py b/src/utils/compactall.py index 6c1da84c..faee36de 100644 --- a/src/utils/compactall.py +++ b/src/utils/compactall.py @@ -1,5 +1,5 @@ import logging -from lib.ursadb import UrsaDb +from lib.ursadb import UrsaDb # type: ignore import time import argparse diff --git a/src/utils/index.py b/src/utils/index.py index 88ab1399..6525a828 100644 --- a/src/utils/index.py +++ b/src/utils/index.py @@ -4,7 +4,7 @@ import json from typing import Set, Iterator, Tuple, List, Optional from pathlib import Path -from lib.ursadb import UrsaDb +from lib.ursadb import UrsaDb # type: ignore from multiprocessing import Pool diff --git a/src/utils/s3index.py b/src/utils/s3index.py index bb92c5ef..b634812b 100644 --- a/src/utils/s3index.py +++ b/src/utils/s3index.py @@ -4,7 +4,7 @@ import shutil from typing import Set, List from pathlib import Path -from lib.ursadb import UrsaDb +from lib.ursadb import UrsaDb # type: ignore from minio import Minio # type: ignore From 80d128a99512d1f42e1e3a15f0b156189b1843b4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Kolankiewicz?= Date: Mon, 30 Sep 2024 10:20:34 +0200 Subject: [PATCH 06/11] Revert "lint: bump mypy version and add ignore types" This reverts commit fa76c06fa85435ec83469b07e0aae92fc6f9ec05. lint: revert todays "fixes" --- .github/workflows/test_code.yml | 4 ++-- CONTRIBUTING.md | 2 +- src/lib/ursadb.py | 4 ++-- src/lib/yaraparse.py | 20 +++++++++---------- src/models/agentgroup.py | 2 +- src/models/configentry.py | 2 +- src/models/job.py | 2 +- src/models/jobagent.py | 2 +- src/models/match.py | 2 +- src/tests/test_ursadb.py | 3 ++- src/tests/test_yaraparse.py | 2 +- .../generate_yaraparse_result_files.py | 4 ++-- src/tests/yararules/test_corpus_yara_rule.py | 4 ++-- src/utils/compactall.py | 2 +- src/utils/index.py | 2 +- src/utils/s3index.py | 2 +- 16 files changed, 30 insertions(+), 29 deletions(-) diff --git a/.github/workflows/test_code.yml b/.github/workflows/test_code.yml index a31074c5..ff1aae71 100644 --- a/.github/workflows/test_code.yml +++ b/.github/workflows/test_code.yml @@ -24,8 +24,8 @@ jobs: uses: actions/setup-python@v1 with: python-version: '3.10' - - name: install mypy==0.981 - run: pip3 install mypy==0.981 + - name: install mypy==0.790 + run: pip3 install mypy==0.790 - name: install requirements run: pip3 install -r requirements.txt - name: run mypy diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index eccec677..0a82fdb3 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -50,7 +50,7 @@ $ prettier --write src/mqueryfront/ - Verify that there are no type errors with [mypy](http://mypy-lang.org/): ```bash -$ pip install mypy==0.981 +$ pip install mypy==0.790 $ mypy src ``` diff --git a/src/lib/ursadb.py b/src/lib/ursadb.py index 786ce8b7..7e97125f 100644 --- a/src/lib/ursadb.py +++ b/src/lib/ursadb.py @@ -48,12 +48,12 @@ def __execute(self, command: str, recv_timeout: int = 2000) -> Json: socket.send_string(command) return json.loads(socket.recv_string()) finally: - socket.close() # type: ignore + socket.close() def query( self, query: str, - taints: List[str] = None, # type: ignore + taints: List[str] = None, dataset: Optional[str] = None, ) -> Json: command = "select " diff --git a/src/lib/yaraparse.py b/src/lib/yaraparse.py index b6377425..44be487d 100644 --- a/src/lib/yaraparse.py +++ b/src/lib/yaraparse.py @@ -377,7 +377,7 @@ def ursify_xor_string(string: PlainString) -> UrsaExpression: def ursify_string(string: String) -> Optional[UrsaExpression]: if string.is_xor: - return ursify_xor_string(string) # type: ignore + return ursify_xor_string(string) elif string.is_plain: return ursify_plain_string( string.pure_text, @@ -389,7 +389,7 @@ def ursify_string(string: String) -> Optional[UrsaExpression]: value_safe = string.pure_text.decode() return ursify_hex(value_safe) elif string.is_regexp: - return ursify_regex_string(string) # type: ignore + return ursify_regex_string(string) return None @@ -429,7 +429,7 @@ def pare_expr( def str_expr( self, condition: StringExpression ) -> Optional[UrsaExpression]: - return ursify_string(self.strings[condition.id]) # type: ignore + return ursify_string(self.strings[condition.id]) def expand_string_wildcard( self, condition: StringWildcardExpression @@ -440,7 +440,7 @@ def expand_string_wildcard( v for k, v in self.strings.items() if re.match(condition_regex, k) ] - ursa_strings = [ursify_string(x) for x in filtered_strings] # type: ignore + ursa_strings = [ursify_string(x) for x in filtered_strings] return [s for s in ursa_strings if s is not None] def expand_set_expression( @@ -463,7 +463,7 @@ def of_expr(self, condition: OfExpression) -> Optional[UrsaExpression]: if type(children) is SetExpression: all_elements = self.expand_set_expression(children) elif type(children) is ThemExpression: - all_elements = [ursify_string(k) for k in self.strings.values()] # type: ignore + all_elements = [ursify_string(k) for k in self.strings.values()] else: raise YaraParseError(f"Unsupported of_expr type: {type(children)}") @@ -524,7 +524,7 @@ def str_count_expr( self, condition: StringCountExpression ) -> Optional[UrsaExpression]: fixed_id = "$" + condition.id[1:] - return ursify_string(self.strings[fixed_id]) # type: ignore + return ursify_string(self.strings[fixed_id]) def int_lit_expr( self, condition: IntLiteralExpression @@ -535,15 +535,15 @@ def int_lit_expr( def str_at_expr( self, condition: StringAtExpression ) -> Optional[UrsaExpression]: - return ursify_string(self.strings[condition.id]) # type: ignore - + return ursify_string(self.strings[condition.id]) + def id_expr(self, condition: IdExpression) -> Optional[UrsaExpression]: return self.rules[condition.symbol.name].parse() def str_in_expr( self, condition: StringInRangeExpression ) -> Optional[UrsaExpression]: - return ursify_string(self.strings[condition.id]) # type: ignore + return ursify_string(self.strings[condition.id]) CONDITION_HANDLERS = { AndExpression: and_expr, @@ -565,7 +565,7 @@ def str_in_expr( def traverse(self, condition) -> Optional[UrsaExpression]: if type(condition) in self.CONDITION_HANDLERS: - return self.CONDITION_HANDLERS[type(condition)](self, condition) # type: ignore + return self.CONDITION_HANDLERS[type(condition)](self, condition) else: print(f"unsupported expression: {type(condition)}") return None diff --git a/src/models/agentgroup.py b/src/models/agentgroup.py index 2325ea25..70fa06a0 100644 --- a/src/models/agentgroup.py +++ b/src/models/agentgroup.py @@ -10,7 +10,7 @@ class AgentGroupBase(SQLModel): active_plugins: List[str] = Field(sa_column=Column(ARRAY(String))) -class AgentGroup(AgentGroupBase, table=True): # type: ignore +class AgentGroup(AgentGroupBase, table=True): """Agent group is a group of processes working on a single file group, with a shared storage, and a single backing ursadb. """ diff --git a/src/models/configentry.py b/src/models/configentry.py index b9cb980b..8e16256b 100644 --- a/src/models/configentry.py +++ b/src/models/configentry.py @@ -2,7 +2,7 @@ from typing import Union -class ConfigEntry(SQLModel, table=True): # type: ignore +class ConfigEntry(SQLModel, table=True): id: Union[int, None] = Field(default=None, primary_key=True) plugin: str key: str diff --git a/src/models/job.py b/src/models/job.py index 2bc8f25d..17db42ad 100644 --- a/src/models/job.py +++ b/src/models/job.py @@ -30,7 +30,7 @@ class JobBase(SQLModel): agents_left: int -class Job(JobBase, table=True): # type: ignore +class Job(JobBase, table=True): """Job object in the database. Internal ID is an implementation detail.""" internal_id: Union[int, None] = Field(default=None, primary_key=True) diff --git a/src/models/jobagent.py b/src/models/jobagent.py index e8a17564..f40313de 100644 --- a/src/models/jobagent.py +++ b/src/models/jobagent.py @@ -6,7 +6,7 @@ from ..models.agentgroup import AgentGroup -class JobAgent(SQLModel, table=True): # type: ignore +class JobAgent(SQLModel, table=True): """Information about job relating to a specific agent group.""" id: Union[int, None] = Field(default=None, primary_key=True) diff --git a/src/models/match.py b/src/models/match.py index e0fcc6cb..23980822 100644 --- a/src/models/match.py +++ b/src/models/match.py @@ -4,7 +4,7 @@ from ..models.job import Job -class Match(SQLModel, table=True): # type: ignore +class Match(SQLModel, table=True): """Represents a file matched to a job, along with a related metadata.""" id: Union[int, None] = Field(default=None, primary_key=True) diff --git a/src/tests/test_ursadb.py b/src/tests/test_ursadb.py index 8de8a999..e628ec80 100644 --- a/src/tests/test_ursadb.py +++ b/src/tests/test_ursadb.py @@ -8,7 +8,8 @@ from typing import Dict, Any sys.path = [".."] + sys.path -from lib.ursadb import UrsaDb # type: ignore +from lib.ursadb import UrsaDb # noqa + class UrsadbTestContext: def __init__(self, socket: zmq.Socket, ursadb: UrsaDb): diff --git a/src/tests/test_yaraparse.py b/src/tests/test_yaraparse.py index 3c9af2b8..c9bda913 100644 --- a/src/tests/test_yaraparse.py +++ b/src/tests/test_yaraparse.py @@ -1,6 +1,6 @@ """Unit tests for yaraparse.""" -from lib.yaraparse import ursify_hex, ursify_plain_string, parse_yara # type: ignore +from lib.yaraparse import ursify_hex, ursify_plain_string, parse_yara import yaramod diff --git a/src/tests/yararules/generate_yaraparse_result_files.py b/src/tests/yararules/generate_yaraparse_result_files.py index 54275526..fa69f11f 100644 --- a/src/tests/yararules/generate_yaraparse_result_files.py +++ b/src/tests/yararules/generate_yaraparse_result_files.py @@ -1,7 +1,7 @@ import argparse import os -from lib.yaraparse import parse_yara # type: ignore -from lib.yaraparse import combine_rules # type: ignore +from lib.yaraparse import parse_yara +from lib.yaraparse import combine_rules current_path = os.path.abspath(os.path.dirname(__file__)) testdir = current_path + "/testdata/" diff --git a/src/tests/yararules/test_corpus_yara_rule.py b/src/tests/yararules/test_corpus_yara_rule.py index 7e513e7a..f34049f1 100644 --- a/src/tests/yararules/test_corpus_yara_rule.py +++ b/src/tests/yararules/test_corpus_yara_rule.py @@ -2,8 +2,8 @@ import unittest from pathlib import Path -from lib.yaraparse import combine_rules # type: ignore -from lib.yaraparse import parse_yara # type: ignore +from lib.yaraparse import combine_rules +from lib.yaraparse import parse_yara testdir = Path(__file__).parent / "testdata" diff --git a/src/utils/compactall.py b/src/utils/compactall.py index faee36de..6c1da84c 100644 --- a/src/utils/compactall.py +++ b/src/utils/compactall.py @@ -1,5 +1,5 @@ import logging -from lib.ursadb import UrsaDb # type: ignore +from lib.ursadb import UrsaDb import time import argparse diff --git a/src/utils/index.py b/src/utils/index.py index 6525a828..88ab1399 100644 --- a/src/utils/index.py +++ b/src/utils/index.py @@ -4,7 +4,7 @@ import json from typing import Set, Iterator, Tuple, List, Optional from pathlib import Path -from lib.ursadb import UrsaDb # type: ignore +from lib.ursadb import UrsaDb from multiprocessing import Pool diff --git a/src/utils/s3index.py b/src/utils/s3index.py index b634812b..bb92c5ef 100644 --- a/src/utils/s3index.py +++ b/src/utils/s3index.py @@ -4,7 +4,7 @@ import shutil from typing import Set, List from pathlib import Path -from lib.ursadb import UrsaDb # type: ignore +from lib.ursadb import UrsaDb from minio import Minio # type: ignore From ce8123d29d45e4b6e294be484330558abba6fc8c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Kolankiewicz?= Date: Mon, 30 Sep 2024 10:21:21 +0200 Subject: [PATCH 07/11] Revert "lint: bump python version" This reverts commit bdd5b1ee585a7e748285e841f59668809f39bf19. --- setup.cfg | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/setup.cfg b/setup.cfg index 2c4a37d2..b4de06f9 100644 --- a/setup.cfg +++ b/setup.cfg @@ -3,7 +3,8 @@ ignore = E501,W503,E203 exclude = mqueryfront/ [mypy] -python_version = 3.10 +python_version = 3.6 +no_site_packages = True [mypy-yaramod.*] ignore_missing_imports = True From dd305ce5edc673e230ec0eb578a8b3afd4143161 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Kolankiewicz?= Date: Mon, 30 Sep 2024 10:21:28 +0200 Subject: [PATCH 08/11] Revert "lint: mypy ignores site-packages directory" This reverts commit 09d7fb5a7bf826cace1a9fdc9e2d79e4b255d84f. --- setup.cfg | 1 - 1 file changed, 1 deletion(-) diff --git a/setup.cfg b/setup.cfg index b4de06f9..f56b362d 100644 --- a/setup.cfg +++ b/setup.cfg @@ -4,7 +4,6 @@ exclude = mqueryfront/ [mypy] python_version = 3.6 -no_site_packages = True [mypy-yaramod.*] ignore_missing_imports = True From 3259dac598be21adc1465cebe67c790ff524480f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Kolankiewicz?= Date: Mon, 30 Sep 2024 12:15:41 +0200 Subject: [PATCH 09/11] lint: revert sql and request bumps, bump mypy version to 1.2.0 --- .github/workflows/test_code.yml | 4 ++-- CONTRIBUTING.md | 2 +- requirements.txt | 4 ++-- setup.cfg | 3 ++- src/app.py | 8 ++++---- src/db.py | 4 ++-- src/lib/ursadb.py | 4 ++-- src/lib/yaraparse.py | 18 +++++++++--------- src/tasks.py | 2 +- src/utils/index.py | 2 +- 10 files changed, 26 insertions(+), 25 deletions(-) diff --git a/.github/workflows/test_code.yml b/.github/workflows/test_code.yml index ff1aae71..eabe404b 100644 --- a/.github/workflows/test_code.yml +++ b/.github/workflows/test_code.yml @@ -24,8 +24,8 @@ jobs: uses: actions/setup-python@v1 with: python-version: '3.10' - - name: install mypy==0.790 - run: pip3 install mypy==0.790 + - name: install mypy==1.2.0 + run: pip3 install mypy==1.2.0 - name: install requirements run: pip3 install -r requirements.txt - name: run mypy diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 0a82fdb3..285d210b 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -50,7 +50,7 @@ $ prettier --write src/mqueryfront/ - Verify that there are no type errors with [mypy](http://mypy-lang.org/): ```bash -$ pip install mypy==0.790 +$ pip install mypy==1.2.0 $ mypy src ``` diff --git a/requirements.txt b/requirements.txt index d5d4ad79..97a1eccb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -19,10 +19,10 @@ pydantic_core==2.23.4 PyJWT[crypto]==2.9.0 pyzmq==26.2.0 redis==5.0.8 -requests==2.32.3 +requests==2.32.2 rq==1.16.2 sniffio==1.3.1 -sqlmodel==0.0.22 +sqlmodel==0.0.11 starlette==0.38.6 typed-config==2.0.3 typing_extensions==4.12.2 diff --git a/setup.cfg b/setup.cfg index f56b362d..1e4c2dc0 100644 --- a/setup.cfg +++ b/setup.cfg @@ -3,7 +3,8 @@ ignore = E501,W503,E203 exclude = mqueryfront/ [mypy] -python_version = 3.6 +python_version = 3.10 +ignore_missing_imports = True [mypy-yaramod.*] ignore_missing_imports = True diff --git a/src/app.py b/src/app.py index eb42bc4f..8bc8fb6d 100644 --- a/src/app.py +++ b/src/app.py @@ -122,7 +122,7 @@ async def current_user(authorization: Optional[str] = Header(None)) -> User: async def add_headers(request: Request, call_next: Callable) -> Response: response = await call_next(request) response.headers["X-Frame-Options"] = "deny" - response.headers["Access-Control-Allow-Origin"] = request.client.host + response.headers["Access-Control-Allow-Origin"] = request.client.host # type: ignore response.headers[ "Access-Control-Allow-Headers" ] = "cache-control,x-requested-with,content-type,authorization" @@ -258,13 +258,13 @@ def backend_status() -> BackendStatusSchema: ursadb_version = status["result"]["ursadb_version"] agents.append( AgentSchema( - name=name, alive=True, tasks=tasks, spec=agent_spec + name=name, alive=True, tasks=tasks, spec=agent_spec # type: ignore ) ) components[f"ursadb ({name})"] = ursadb_version except Again: agents.append( - AgentSchema(name=name, alive=False, tasks=[], spec=agent_spec) + AgentSchema(name=name, alive=False, tasks=[], spec=agent_spec) # type: ignore ) components[f"ursadb ({name})"] = "unknown" @@ -534,7 +534,7 @@ def job_statuses(user: User = Depends(current_user)) -> JobsSchema: if "can_list_all_queries" in get_user_roles(user): username_filter = None jobs = db.get_valid_jobs(username_filter) - return JobsSchema(jobs=jobs) + return JobsSchema(jobs=jobs) # type: ignore @app.delete( diff --git a/src/db.py b/src/db.py index 47036d0a..43395233 100644 --- a/src/db.py +++ b/src/db.py @@ -267,7 +267,7 @@ def create_search_task( datasets_left=0, total_datasets=0, taints=taints, - ) + ) # type: ignore session.add(obj) session.commit() @@ -287,7 +287,7 @@ def get_job_matches( query = query.limit(limit) matches = session.exec(query).all() - return MatchesSchema(job=job, matches=matches) + return MatchesSchema(job=job, matches=matches) # type: ignore def update_job_files(self, job: JobId, total_files: int) -> int: """Add total_files to the specified job, and return a new total.""" diff --git a/src/lib/ursadb.py b/src/lib/ursadb.py index 7e97125f..c7fdf412 100644 --- a/src/lib/ursadb.py +++ b/src/lib/ursadb.py @@ -48,12 +48,12 @@ def __execute(self, command: str, recv_timeout: int = 2000) -> Json: socket.send_string(command) return json.loads(socket.recv_string()) finally: - socket.close() + socket.close() # type: ignore def query( self, query: str, - taints: List[str] = None, + taints: List[str] | None = None, dataset: Optional[str] = None, ) -> Json: command = "select " diff --git a/src/lib/yaraparse.py b/src/lib/yaraparse.py index 44be487d..99336104 100644 --- a/src/lib/yaraparse.py +++ b/src/lib/yaraparse.py @@ -377,7 +377,7 @@ def ursify_xor_string(string: PlainString) -> UrsaExpression: def ursify_string(string: String) -> Optional[UrsaExpression]: if string.is_xor: - return ursify_xor_string(string) + return ursify_xor_string(string) # type: ignore elif string.is_plain: return ursify_plain_string( string.pure_text, @@ -389,7 +389,7 @@ def ursify_string(string: String) -> Optional[UrsaExpression]: value_safe = string.pure_text.decode() return ursify_hex(value_safe) elif string.is_regexp: - return ursify_regex_string(string) + return ursify_regex_string(string) # type: ignore return None @@ -429,7 +429,7 @@ def pare_expr( def str_expr( self, condition: StringExpression ) -> Optional[UrsaExpression]: - return ursify_string(self.strings[condition.id]) + return ursify_string(self.strings[condition.id]) # type: ignore def expand_string_wildcard( self, condition: StringWildcardExpression @@ -440,7 +440,7 @@ def expand_string_wildcard( v for k, v in self.strings.items() if re.match(condition_regex, k) ] - ursa_strings = [ursify_string(x) for x in filtered_strings] + ursa_strings = [ursify_string(x) for x in filtered_strings] # type: ignore return [s for s in ursa_strings if s is not None] def expand_set_expression( @@ -463,7 +463,7 @@ def of_expr(self, condition: OfExpression) -> Optional[UrsaExpression]: if type(children) is SetExpression: all_elements = self.expand_set_expression(children) elif type(children) is ThemExpression: - all_elements = [ursify_string(k) for k in self.strings.values()] + all_elements = [ursify_string(k) for k in self.strings.values()] # type: ignore else: raise YaraParseError(f"Unsupported of_expr type: {type(children)}") @@ -524,7 +524,7 @@ def str_count_expr( self, condition: StringCountExpression ) -> Optional[UrsaExpression]: fixed_id = "$" + condition.id[1:] - return ursify_string(self.strings[fixed_id]) + return ursify_string(self.strings[fixed_id]) # type: ignore def int_lit_expr( self, condition: IntLiteralExpression @@ -535,7 +535,7 @@ def int_lit_expr( def str_at_expr( self, condition: StringAtExpression ) -> Optional[UrsaExpression]: - return ursify_string(self.strings[condition.id]) + return ursify_string(self.strings[condition.id]) # type: ignore def id_expr(self, condition: IdExpression) -> Optional[UrsaExpression]: return self.rules[condition.symbol.name].parse() @@ -543,7 +543,7 @@ def id_expr(self, condition: IdExpression) -> Optional[UrsaExpression]: def str_in_expr( self, condition: StringInRangeExpression ) -> Optional[UrsaExpression]: - return ursify_string(self.strings[condition.id]) + return ursify_string(self.strings[condition.id]) # type: ignore CONDITION_HANDLERS = { AndExpression: and_expr, @@ -565,7 +565,7 @@ def str_in_expr( def traverse(self, condition) -> Optional[UrsaExpression]: if type(condition) in self.CONDITION_HANDLERS: - return self.CONDITION_HANDLERS[type(condition)](self, condition) + return self.CONDITION_HANDLERS[type(condition)](self, condition) # type: ignore else: print(f"unsupported expression: {type(condition)}") return None diff --git a/src/tasks.py b/src/tasks.py index edb9a426..5e3e105c 100644 --- a/src/tasks.py +++ b/src/tasks.py @@ -165,7 +165,7 @@ def make_agent(group_override: Optional[str] = None): if group_override is not None: group_id = group_override else: - group_id = get_current_job().origin + group_id = get_current_job().origin # type: ignore return Agent(group_id) diff --git a/src/utils/index.py b/src/utils/index.py index 88ab1399..3a211f95 100644 --- a/src/utils/index.py +++ b/src/utils/index.py @@ -287,7 +287,7 @@ def main() -> None: path_mount = args.path path = Path(args.path) - if not path.exists: + if not path.exists(): logging.error("Path (--path) %s does not exist.", args.path) return From 3dc9e3096f4605931c381f15f63b548ecc794458 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Kolankiewicz?= Date: Mon, 30 Sep 2024 14:00:28 +0200 Subject: [PATCH 10/11] fix: types-requests added --- requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements.txt b/requirements.txt index 97a1eccb..305158fe 100644 --- a/requirements.txt +++ b/requirements.txt @@ -25,6 +25,7 @@ sniffio==1.3.1 sqlmodel==0.0.11 starlette==0.38.6 typed-config==2.0.3 +types-requests==2.32.0.20240914 typing_extensions==4.12.2 urllib3==2.2.3 uvicorn==0.30.6 From 49614901058bfaefe4ce5a45606e7c99e6ae06b3 Mon Sep 17 00:00:00 2001 From: msm Date: Wed, 2 Oct 2024 16:03:55 +0200 Subject: [PATCH 11/11] Don't ignore type errors --- .github/workflows/test_code.yml | 8 ++++++-- requirements.txt | 2 +- setup.cfg | 1 - src/app.py | 6 +++--- src/db.py | 4 ++-- src/lib/ursadb.py | 2 +- src/lib/yaraparse.py | 27 +++++++++++++-------------- src/models/agentgroup.py | 10 ++-------- src/models/job.py | 12 +++--------- src/schema.py | 4 ++-- 10 files changed, 33 insertions(+), 43 deletions(-) diff --git a/.github/workflows/test_code.yml b/.github/workflows/test_code.yml index eabe404b..0c711859 100644 --- a/.github/workflows/test_code.yml +++ b/.github/workflows/test_code.yml @@ -28,8 +28,12 @@ jobs: run: pip3 install mypy==1.2.0 - name: install requirements run: pip3 install -r requirements.txt - - name: run mypy - run: mypy src + - name: run mypy on main files + run: mypy src/app.py src/daemon.py + - name: run mypy on tests + run: MYPYPATH=src/ mypy src/tests/ + - name: run mypy on utils + run: MYPYPATH=src/ mypy src/utils/ test_python_style: name: python flake8 runs-on: ubuntu-latest diff --git a/requirements.txt b/requirements.txt index 305158fe..27da0663 100644 --- a/requirements.txt +++ b/requirements.txt @@ -31,4 +31,4 @@ urllib3==2.2.3 uvicorn==0.30.6 wrapt==1.16.0 yara-python==4.5.1 -yaramod==3.23.0 \ No newline at end of file +yaramod==3.23.0 diff --git a/setup.cfg b/setup.cfg index 1e4c2dc0..2c4a37d2 100644 --- a/setup.cfg +++ b/setup.cfg @@ -4,7 +4,6 @@ exclude = mqueryfront/ [mypy] python_version = 3.10 -ignore_missing_imports = True [mypy-yaramod.*] ignore_missing_imports = True diff --git a/src/app.py b/src/app.py index 8bc8fb6d..c86d70c9 100644 --- a/src/app.py +++ b/src/app.py @@ -258,13 +258,13 @@ def backend_status() -> BackendStatusSchema: ursadb_version = status["result"]["ursadb_version"] agents.append( AgentSchema( - name=name, alive=True, tasks=tasks, spec=agent_spec # type: ignore + name=name, alive=True, tasks=tasks, spec=agent_spec ) ) components[f"ursadb ({name})"] = ursadb_version except Again: agents.append( - AgentSchema(name=name, alive=False, tasks=[], spec=agent_spec) # type: ignore + AgentSchema(name=name, alive=False, tasks=[], spec=agent_spec) ) components[f"ursadb ({name})"] = "unknown" @@ -534,7 +534,7 @@ def job_statuses(user: User = Depends(current_user)) -> JobsSchema: if "can_list_all_queries" in get_user_roles(user): username_filter = None jobs = db.get_valid_jobs(username_filter) - return JobsSchema(jobs=jobs) # type: ignore + return JobsSchema(jobs=jobs) @app.delete( diff --git a/src/db.py b/src/db.py index 43395233..47036d0a 100644 --- a/src/db.py +++ b/src/db.py @@ -267,7 +267,7 @@ def create_search_task( datasets_left=0, total_datasets=0, taints=taints, - ) # type: ignore + ) session.add(obj) session.commit() @@ -287,7 +287,7 @@ def get_job_matches( query = query.limit(limit) matches = session.exec(query).all() - return MatchesSchema(job=job, matches=matches) # type: ignore + return MatchesSchema(job=job, matches=matches) def update_job_files(self, job: JobId, total_files: int) -> int: """Add total_files to the specified job, and return a new total.""" diff --git a/src/lib/ursadb.py b/src/lib/ursadb.py index c7fdf412..7e995b9f 100644 --- a/src/lib/ursadb.py +++ b/src/lib/ursadb.py @@ -48,7 +48,7 @@ def __execute(self, command: str, recv_timeout: int = 2000) -> Json: socket.send_string(command) return json.loads(socket.recv_string()) finally: - socket.close() # type: ignore + socket.close() def query( self, diff --git a/src/lib/yaraparse.py b/src/lib/yaraparse.py index 99336104..b353520b 100644 --- a/src/lib/yaraparse.py +++ b/src/lib/yaraparse.py @@ -1,7 +1,7 @@ import argparse import itertools import re -from typing import Any, Dict, List, Match, Optional +from typing import Any, Dict, List, Match, Optional, cast, Callable from yaramod import ( # type: ignore AllExpression, @@ -17,7 +17,6 @@ OfExpression, OrExpression, ParenthesesExpression, - PlainString, Regexp, RegexpConcat, RegexpGroup, @@ -358,7 +357,7 @@ def ursify_plain_string( return ursa_ascii -def ursify_xor_string(string: PlainString) -> UrsaExpression: +def ursify_xor_string(string: String) -> UrsaExpression: text_ascii = string.pure_text xored_strings: List[UrsaExpression] = [] @@ -377,7 +376,7 @@ def ursify_xor_string(string: PlainString) -> UrsaExpression: def ursify_string(string: String) -> Optional[UrsaExpression]: if string.is_xor: - return ursify_xor_string(string) # type: ignore + return ursify_xor_string(string) elif string.is_plain: return ursify_plain_string( string.pure_text, @@ -389,14 +388,14 @@ def ursify_string(string: String) -> Optional[UrsaExpression]: value_safe = string.pure_text.decode() return ursify_hex(value_safe) elif string.is_regexp: - return ursify_regex_string(string) # type: ignore + return ursify_regex_string(cast(Regexp, string)) return None class RuleParseEngine: def __init__( - self, strings: Dict[str, str], rules: Dict[str, YaraRuleData] + self, strings: Dict[str, String], rules: Dict[str, YaraRuleData] ) -> None: self.strings = strings self.rules = rules @@ -429,7 +428,7 @@ def pare_expr( def str_expr( self, condition: StringExpression ) -> Optional[UrsaExpression]: - return ursify_string(self.strings[condition.id]) # type: ignore + return ursify_string(self.strings[condition.id]) def expand_string_wildcard( self, condition: StringWildcardExpression @@ -440,7 +439,7 @@ def expand_string_wildcard( v for k, v in self.strings.items() if re.match(condition_regex, k) ] - ursa_strings = [ursify_string(x) for x in filtered_strings] # type: ignore + ursa_strings = [ursify_string(x) for x in filtered_strings] return [s for s in ursa_strings if s is not None] def expand_set_expression( @@ -463,7 +462,7 @@ def of_expr(self, condition: OfExpression) -> Optional[UrsaExpression]: if type(children) is SetExpression: all_elements = self.expand_set_expression(children) elif type(children) is ThemExpression: - all_elements = [ursify_string(k) for k in self.strings.values()] # type: ignore + all_elements = [ursify_string(k) for k in self.strings.values()] else: raise YaraParseError(f"Unsupported of_expr type: {type(children)}") @@ -524,7 +523,7 @@ def str_count_expr( self, condition: StringCountExpression ) -> Optional[UrsaExpression]: fixed_id = "$" + condition.id[1:] - return ursify_string(self.strings[fixed_id]) # type: ignore + return ursify_string(self.strings[fixed_id]) def int_lit_expr( self, condition: IntLiteralExpression @@ -535,7 +534,7 @@ def int_lit_expr( def str_at_expr( self, condition: StringAtExpression ) -> Optional[UrsaExpression]: - return ursify_string(self.strings[condition.id]) # type: ignore + return ursify_string(self.strings[condition.id]) def id_expr(self, condition: IdExpression) -> Optional[UrsaExpression]: return self.rules[condition.symbol.name].parse() @@ -543,9 +542,9 @@ def id_expr(self, condition: IdExpression) -> Optional[UrsaExpression]: def str_in_expr( self, condition: StringInRangeExpression ) -> Optional[UrsaExpression]: - return ursify_string(self.strings[condition.id]) # type: ignore + return ursify_string(self.strings[condition.id]) - CONDITION_HANDLERS = { + CONDITION_HANDLERS: Dict[type, Callable] = { AndExpression: and_expr, OrExpression: or_expr, ParenthesesExpression: pare_expr, @@ -565,7 +564,7 @@ def str_in_expr( def traverse(self, condition) -> Optional[UrsaExpression]: if type(condition) in self.CONDITION_HANDLERS: - return self.CONDITION_HANDLERS[type(condition)](self, condition) # type: ignore + return self.CONDITION_HANDLERS[type(condition)](self, condition) else: print(f"unsupported expression: {type(condition)}") return None diff --git a/src/models/agentgroup.py b/src/models/agentgroup.py index 70fa06a0..b5bfa39e 100644 --- a/src/models/agentgroup.py +++ b/src/models/agentgroup.py @@ -3,23 +3,17 @@ from ..models.jobagent import JobAgent -class AgentGroupBase(SQLModel): +class AgentGroupView(SQLModel): name: str ursadb_url: str plugins_spec: Dict[str, Dict[str, str]] = Field(sa_column=Column(JSON)) active_plugins: List[str] = Field(sa_column=Column(ARRAY(String))) -class AgentGroup(AgentGroupBase, table=True): +class AgentGroup(AgentGroupView, table=True): """Agent group is a group of processes working on a single file group, with a shared storage, and a single backing ursadb. """ id: Union[int, None] = Field(default=None, primary_key=True) jobs: List["JobAgent"] = Relationship(back_populates="agent") - - -class AgentGroupView(AgentGroupBase): - """Pydantic model used in the public API.""" - - pass diff --git a/src/models/job.py b/src/models/job.py index 17db42ad..4fa93abd 100644 --- a/src/models/job.py +++ b/src/models/job.py @@ -6,8 +6,8 @@ from ..models.jobagent import JobAgent -class JobBase(SQLModel): - """Base class for entities related to mquery jobs.""" +class JobView(SQLModel): + """Public fields of mquery jobs.""" id: str status: str @@ -30,16 +30,10 @@ class JobBase(SQLModel): agents_left: int -class Job(JobBase, table=True): +class Job(JobView, table=True): """Job object in the database. Internal ID is an implementation detail.""" internal_id: Union[int, None] = Field(default=None, primary_key=True) matches: List["Match"] = Relationship(back_populates="job") agents: List["JobAgent"] = Relationship(back_populates="job") - - -class JobView(JobBase): - """Pydantic model used in the public API.""" - - pass diff --git a/src/schema.py b/src/schema.py index a4b285fb..5b2ab8ff 100644 --- a/src/schema.py +++ b/src/schema.py @@ -1,12 +1,12 @@ from enum import Enum -from typing import List, Dict, Optional +from typing import List, Dict, Optional, Sequence from pydantic import BaseModel, Field # type: ignore from .models.job import JobView from .models.agentgroup import AgentGroupView class JobsSchema(BaseModel): - jobs: List[JobView] + jobs: Sequence[JobView] class ConfigSchema(BaseModel):