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..f6b8a7f5 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, @@ -358,7 +358,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 +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,14 +389,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 +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,7 +535,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 +543,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 +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 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):