diff --git a/ci/apiv2/test_agent.py b/ci/apiv2/test_agent.py index 83e619091..c044b573a 100644 --- a/ci/apiv2/test_agent.py +++ b/ci/apiv2/test_agent.py @@ -2,6 +2,7 @@ import test_user from hashtopolis import Agent, Config, Helper from hashtopolis import HashtopolisError +from hashtopolis_agent import ProcessState from utils import BaseTest @@ -110,3 +111,23 @@ def test_hide_ip_info(self): def test_acl(self): model_obj = self.create_test_object() self._test_acl_list(model_obj, {'permAgentRead': True}) + + def test_active_chunk(self): + dummy_agent, agent, _, task = self.create_agent_with_task().values() + dummy_agent.get_chunk() + dummy_agent.send_process(progress=50, state=ProcessState.RUNNING) + agent_resp = Agent.objects.get(pk=agent.id) + chunks = agent_resp._Model__relationships['chunks'].get('data', []) + active_chunk_id = next(iter(chunks), {}).get('id', None) + + self.assertEqual(dummy_agent.chunk['chunkId'], active_chunk_id, "Active chunk is reported incorrectly") + + dummy_agent.get_chunk() + # To simulate a stop instruction (e.g. after a hashlist has been completed by other agents), + # report completeness on current chunk, but leave state as RUNNING. + dummy_agent.send_process(progress=100, state=ProcessState.RUNNING) + helper = Helper() + result = helper.unassign_agent(agent=agent) + agent_resp = Agent.objects.get(pk=agent.id) + + self.assertNotIn('data', agent_resp._Model__relationships['chunks'], "Chunks of a completed hashlist should not be returned as active") \ No newline at end of file diff --git a/src/inc/apiv2/model/AgentAPI.php b/src/inc/apiv2/model/AgentAPI.php index 0e83690f1..7a5f6bde5 100644 --- a/src/inc/apiv2/model/AgentAPI.php +++ b/src/inc/apiv2/model/AgentAPI.php @@ -6,7 +6,6 @@ use Hashtopolis\dba\AbstractModel; use Hashtopolis\inc\utils\AccessUtils; use Hashtopolis\inc\utils\AgentUtils; -use Hashtopolis\inc\defines\DHashcatStatus; use Hashtopolis\dba\ContainFilter; use Hashtopolis\dba\ExistsFilter; use Hashtopolis\dba\Factory; @@ -18,7 +17,6 @@ use Hashtopolis\dba\models\AgentStat; use Hashtopolis\dba\models\Assignment; use Hashtopolis\dba\models\Chunk; -use Hashtopolis\dba\QueryFilter; use Hashtopolis\dba\models\Task; use Hashtopolis\dba\models\User; use Hashtopolis\inc\apiv2\common\AbstractModelAPI; @@ -92,15 +90,10 @@ protected function getAggregateCrackingTime(AbstractModel $object): int { */ function aggregateData(AbstractModel $object, array &$includedData = [], ?array $aggregateFieldsets = null): array { $agentId = $object->getId(); - $qFs = []; - $qFs[] = new QueryFilter(Chunk::AGENT_ID, $agentId, "="); - $qFs[] = new QueryFilter(Chunk::STATE, DHashcatStatus::RUNNING, "="); - - $active_chunk = Factory::getChunkFactory()->filter([Factory::FILTER => $qFs], true); - if ($active_chunk !== NULL) { + $active_chunk = AgentUtils::getActiveChunk($agentId); + if ($active_chunk !== null) { $includedData["chunks"][$agentId] = [$active_chunk]; } - return parent::aggregateData($object, $includedData, $aggregateFieldsets); } diff --git a/src/inc/apiv2/model/AgentAssignmentAPI.php b/src/inc/apiv2/model/AgentAssignmentAPI.php index b3f8e30c8..8b4412480 100644 --- a/src/inc/apiv2/model/AgentAssignmentAPI.php +++ b/src/inc/apiv2/model/AgentAssignmentAPI.php @@ -4,10 +4,6 @@ use Exception; use Hashtopolis\dba\AbstractModel; -use Hashtopolis\dba\models\Chunk; -use Hashtopolis\dba\QueryFilter; -use Hashtopolis\inc\defines\DConfig; -use Hashtopolis\inc\SConfig; use Hashtopolis\inc\utils\AccessUtils; use Hashtopolis\inc\utils\AgentUtils; use Hashtopolis\inc\utils\AssignmentUtils; @@ -155,12 +151,8 @@ protected function getAggregateCurrentSpeed(AbstractModel $object): int { * @throws Exception */ protected function getAggregateCurrentChunkId(AbstractModel $object): ?int { - $qF1 = new QueryFilter(Chunk::TASK_ID, $object->getTaskId(), "="); - $qF2 = new QueryFilter(Chunk::AGENT_ID, $object->getAgentId(), "="); - $qF3 = new QueryFilter(Chunk::SOLVE_TIME, time() - SConfig::getInstance()->getVal(DConfig::CHUNK_TIMEOUT), ">"); - $qF4 = new QueryFilter(Chunk::PROGRESS, 10000, "<"); - $chunk = Factory::getChunkFactory()->filter([Factory::FILTER => array_filter([$qF1, $qF2, $qF3, $qF4])], true); - return $chunk?->getId(); + $active_chunk = AgentUtils::getActiveChunk($object->getAgentId(), $object->getTaskId()); + return $active_chunk?->getId(); } /** diff --git a/src/inc/utils/AgentUtils.php b/src/inc/utils/AgentUtils.php index 8fb9b7d0d..bfd1f2fdb 100644 --- a/src/inc/utils/AgentUtils.php +++ b/src/inc/utils/AgentUtils.php @@ -27,6 +27,7 @@ use Hashtopolis\inc\apiv2\error\HttpError; use Hashtopolis\inc\defines\DAgentStatsType; use Hashtopolis\inc\defines\DConfig; +use Hashtopolis\inc\defines\DHashcatStatus; use Hashtopolis\inc\defines\DLogEntry; use Hashtopolis\inc\defines\DLogEntryIssuer; use Hashtopolis\inc\defines\DNotificationObjectType; @@ -644,4 +645,24 @@ public static function getAggregateCracked(int $agentId, ?int $taskId = null): i $results = Factory::getChunkFactory()->multicolAggregationFilter([Factory::FILTER => array_filter([$qF1, $qF2])], [$agg1]); return (int)($results[$agg1->getName()] ?? 0); } + + /** + * Get the active chunk being worked on by an agent or + * (if task ID is specified) by an agent on a specific task. + * + * @param int $agentId + * @param int|null $taskId + * @return Chunk|null + * @throws Exception + */ + public static function getActiveChunk(int $agentId, ?int $taskId = null): ?Chunk { + $qFs = []; + $qFs[] = new QueryFilter(Chunk::AGENT_ID, $agentId, "="); + $qFs[] = $taskId !== null ? new QueryFilter(Chunk::TASK_ID, $taskId, "=") : null; + $qFs[] = new QueryFilter(Chunk::STATE, DHashcatStatus::RUNNING, "="); + $qFs[] = new QueryFilter(Chunk::SOLVE_TIME, time() - SConfig::getInstance()->getVal(DConfig::CHUNK_TIMEOUT), ">"); + $qFs[] = new QueryFilter(Chunk::PROGRESS, 10000, "<"); + $oF = new OrderFilter(Chunk::SOLVE_TIME, "DESC"); + return Factory::getChunkFactory()->filter([Factory::FILTER => array_filter($qFs), Factory::ORDER => $oF], true); + } }