feat: Enhance press-agent communication

.cspell.json

@@ -794,4 +794,4 @@
         "libgnutls",
         "UDF"
     ]
-}
+}

press/agent.py

@@ -12,6 +12,7 @@
 
 import frappe
 import frappe.utils
+import jwt
 import requests
 from frappe.utils.password import get_decrypted_password
 from requests.exceptions import HTTPError
@@ -949,13 +950,62 @@ def _make_req(self, method, path, data, files, agent_job_id):
 			return requests.request(method, url, headers=headers, files=file_objects, verify=verify)
 		return requests.request(method, url, headers=headers, json=data, verify=verify, timeout=(10, 30))
 
+	def get_secret(self):
+		key = "agent_auth_secret"
+
+		secret = frappe.cache().get_value(key)
+		if not secret:
+			press_settings = frappe.get_single("Press Settings")
+			secret = press_settings.get_password("secret")
+
+			if not secret:
+				raise ValueError("Agent auth secret not configured")
+
+			frappe.cache().set_value(key, secret)
+
+		return secret
+
+	def _verify_request_token(self, token: str):
+		secret = self.get_secret()
+
+		try:
+			payload = jwt.decode(
+				token,
+				secret,
+				algorithms=["HS256"],
+				options={
+					"require": ["exp", "server", "jti"],
+				},
+			)
+
+		except jwt.ExpiredSignatureError as err:
+			raise ValueError("Token expired") from err
+
+		except jwt.InvalidTokenError as err:
+			raise ValueError("Invalid token") from err
+
+		if payload["server"] != self.server:
+			raise ValueError("Invalid server")
+
+		return True
+
+	def extract_and_verify_token(self, token):
+		if not token:
+			frappe.throw("Unsigned request from agent", frappe.PermissionError)
+
+		try:
+			self._verify_request_token(token=token)
+		except ValueError:
+			frappe.throw_permission_error()
+
 	def request(self, method, path, data=None, files=None, agent_job=None, raises=True):
 		self.raise_if_past_requests_have_failed()
 		response = json_response = None
 		try:
 			agent_job_id = agent_job.name if agent_job else None
 			response = self._make_req(method, path, data, files, agent_job_id)
 			json_response = response.json()
+
 			if raises and response.status_code >= 400:
 				output = "\n\n".join([json_response.get("output", ""), json_response.get("traceback", "")])
 				if output == "\n\n":
@@ -1026,6 +1076,7 @@ def raw_request(self, method, path, data=None, raises=True, timeout=None):
 		timeout = timeout or (10, 30)
 		response = requests.request(method, url, headers=headers, json=data, timeout=timeout)
 		json_response = response.json()
+
 		if raises:
 			response.raise_for_status()
 		return json_response
@@ -1265,6 +1316,12 @@ def fetch_bench_status(self, bench):
 	def get_snapshot(self, bench: str):
 		return self.get(f"process-snapshot/{bench}")
 
+	def enable_feature_flag(self):
+		return self.post("server/feature/enable")
+
+	def disable_feature_flag(self):
+		return self.post("server/feature/disable")
+
 	def run_after_migrate_steps(self, site):
 		data = {
 			"admin_password": site.get_password("admin_password"),

press/api/agent_auth.py

@@ -0,0 +1,61 @@
+import base64
+import json
+from typing import TYPE_CHECKING
+
+import frappe
+from frappe.rate_limiter import rate_limit
+
+from press.agent import Agent
+
+if TYPE_CHECKING:
+	from press.press.doctype.server.server import BaseServer
+
+
+def extract_server_from_token(token: str):
+	try:
+		parts = token.split(".")
+
+		if len(parts) != 3:
+			return None
+
+		payload_b64 = parts[1]
+
+		# fix padding
+		payload_b64 += "=" * (-len(payload_b64) % 4)
+
+		payload = json.loads(base64.urlsafe_b64decode(payload_b64))
+
+		return payload.get("server"), payload.get("server_type")
+
+	except Exception:
+		return None
+
+
+def verify_agent():
+	agent_token = frappe.request.headers.get("X-Agent-Token")
+
+	if not agent_token:
+		frappe.throw_permission_error()
+
+	token_data = extract_server_from_token(agent_token)
+
+	if not token_data:
+		frappe.throw_permission_error()
+
+	server, server_type = token_data
+
+	agent = Agent(server)
+	agent.extract_and_verify_token(agent_token)
+
+	return server, server_type
+
+
+@frappe.whitelist(allow_guest=True)
+@rate_limit(limit=10, seconds=60)
+def regenerate_token():
+	server, server_type = verify_agent()
+
+	doc: BaseServer = frappe.get_doc(server_type, server)
+
+	secret = doc._generate_secret()
+	return doc.sign_agent_token(secret)

press/api/callbacks.py

@@ -3,12 +3,15 @@
 from __future__ import annotations
 
 import ipaddress
+import json
+from typing import Any
 
 import frappe
 from frappe.rate_limiter import rate_limit
 
 from press.agent import Agent
-from press.press.doctype.agent_job.agent_job import handle_polled_job
+from press.api.agent_auth import verify_agent
+from press.press.doctype.agent_job.agent_job import handle_polled_job, retry_undelivered_jobs
 from press.utils import log_error
 
 
@@ -113,8 +116,57 @@ def callback(job_id: str | None = None):
 	if not server:
 		frappe.throw("Not permitted", frappe.ValidationError)
 
+	verify_agent()
+
 	job = verify_job_id(server, job_id)
 	if not job:
 		frappe.throw("Invalid Job Id", frappe.ValidationError)
 
 	frappe.enqueue(handle_job_updates, server=server, job_identifier=job_id)
+
+
+@frappe.whitelist(allow_guest=True)
+@rate_limit(limit=500, seconds=60)
+def update_job(job: str) -> None:
+	if not job:
+		return
+
+	server, _ = verify_agent()
+
+	parsed_job: dict[str, Any] = json.loads(job)
+
+	job_doc = frappe.get_value(
+		"Agent Job",
+		fieldname=[
+			"name",
+			"job_id",
+			"status",
+			"callback_failure_count",
+			"job_type",
+		],
+		filters={"job_id": parsed_job["id"], "server": server},
+		as_dict=True,
+	)
+
+	if not job_doc:
+		return
+
+	handle_polled_job(
+		polled_job=parsed_job,
+		job=job_doc,
+		raise_callback_exception=True,
+	)
+
+
+@frappe.whitelist(allow_guest=True)
+def retry_undelivered():
+	server, server_type = verify_agent()
+
+	server_obj = frappe._dict(
+		{
+			"server": server,
+			"server_type": server_type,
+		}
+	)
+
+	retry_undelivered_jobs(server_obj, use_exponential_backoff=False, use_queue_protection=True)

press/api/server.py

@@ -14,6 +14,7 @@
 from frappe.utils.password import get_decrypted_password
 
 from press.api.account import is_limits_exceeded
+from press.api.agent_auth import verify_agent
 from press.api.analytics import auto_timespan_timegrain, get_rounded_boundaries, get_rounded_boundary
 from press.api.bench import all as all_benches
 from press.api.site import protected
@@ -987,20 +988,16 @@ def rename(name, title):
 
 
 @frappe.whitelist(allow_guest=True)
-def benches_are_idle(server: str, access_token: str) -> None:
+def benches_are_idle() -> None:
 	"""Shut down the secondary server if all benches are idle.
 
 	This function is only triggered by secondary servers:
 	https://github.com/frappe/agent/pull/346/files#diff-7355d9c50cadfa3f4c74fc77a4ad8ab08e4da8f6c3326bbf9b0de0f00a0aa0daR87-R93
 	"""
-	from passlib.hash import pbkdf2_sha256 as pbkdf2
 
-	server_doc = frappe.get_cached_doc("Server", server)
-	agent_password = server_doc.get_password("agent_password")
 	current_user = frappe.session.user
 
-	if not pbkdf2.verify(agent_password, access_token):
-		return
+	server, _ = verify_agent()
 
 	primary_server, is_server_scaled_up = frappe.db.get_value(
 		"Server", {"secondary_server": server}, ["name", "scaled_up"]

press/api/tests/test_agent_auth.py

@@ -0,0 +1,47 @@
+from unittest import TestCase
+from unittest.mock import Mock, patch
+
+import frappe
+
+from press.api.agent_auth import verify_agent
+
+
+class TestAgentAuth(TestCase):
+	def tearDown(self):
+		frappe.db.rollback()
+
+	def test_verify_agent_throws_without_token(self):
+		frappe.local.request = frappe._dict({"headers": {}})
+
+		self.assertRaises(
+			frappe.PermissionError,
+			verify_agent,
+		)
+
+	@patch("press.api.agent_auth.extract_server_from_token")
+	@patch("press.api.agent_auth.Agent")
+	def test_verify_agent_calls_extract_and_verify_token(
+		self,
+		mock_agent,
+		mock_extract_server,
+	):
+		mock_extract_server.return_value = (
+			"test-server",
+			"Server",
+		)
+
+		mock_instance = Mock()
+		mock_agent.return_value = mock_instance
+
+		frappe.local.request = frappe._dict({"headers": {"X-Agent-Token": "test-token"}})
+
+		server, server_type = verify_agent()
+
+		self.assertEqual(server, "test-server")
+		self.assertEqual(server_type, "Server")
+
+		mock_extract_server.assert_called_once_with("test-token")
+
+		mock_agent.assert_called_once_with("test-server")
+
+		mock_instance.extract_and_verify_token.assert_called_once_with("test-token")

press/api/tests/test_callbacks.py

@@ -0,0 +1,60 @@
+import json
+from unittest import TestCase
+from unittest.mock import patch
+
+import frappe
+
+from press.api.callbacks import (
+	retry_undelivered,
+	update_job,
+)
+
+
+class TestCallbacks(TestCase):
+	def tearDown(self):
+		frappe.db.rollback()
+
+	def test_update_job_returns_when_job_missing(self):
+		self.assertIsNone(update_job(job=None))
+
+	@patch("press.api.callbacks.handle_polled_job")
+	@patch("press.api.callbacks.verify_agent")
+	@patch("frappe.get_value")
+	def test_update_job_calls_handle_polled_job(
+		self,
+		mock_get_value,
+		mock_verify,
+		mock_handle_polled_job,
+	):
+		mock_verify.return_value = (
+			"test-server",
+			"Server",
+		)
+
+		mock_get_value.return_value = {
+			"name": "job-1",
+			"job_id": "123",
+			"status": "Running",
+			"callback_failure_count": 0,
+			"job_type": "Deploy",
+		}
+
+		update_job(job=json.dumps({"id": "123"}))
+
+		mock_handle_polled_job.assert_called_once()
+
+	@patch("press.api.callbacks.retry_undelivered_jobs")
+	@patch("press.api.callbacks.verify_agent")
+	def test_retry_undelivered(
+		self,
+		mock_verify,
+		mock_retry_jobs,
+	):
+		mock_verify.return_value = (
+			"test-server",
+			"Server",
+		)
+
+		retry_undelivered()
+
+		mock_retry_jobs.assert_called_once()

press/playbooks/roles/agent/tasks/main.yml

@@ -21,7 +21,7 @@
 - name: Generate Agent Configuration File
   become: yes
   become_user: frappe
-  command: '/home/frappe/agent/env/bin/agent setup config --name {{ server }} --workers {{ workers }} {% if proxy_ip is defined and proxy_ip is truthy %}--proxy-ip {{ proxy_ip }}{% endif %} {% if agent_sentry_dsn is defined and agent_sentry_dsn is truthy %}--sentry-dsn {{ agent_sentry_dsn }}{% endif %}'
+  command: '/home/frappe/agent/env/bin/agent setup config --name {{ server }} --workers {{ workers }} {% if proxy_ip is defined and proxy_ip is truthy %}--proxy-ip {{ proxy_ip }}{% endif %} {% if agent_sentry_dsn is defined and agent_sentry_dsn is truthy %}--sentry-dsn {{ agent_sentry_dsn }}{% endif %} {% if agent_token is defined and agent_token is truthy %}--agent-token {{ agent_token }}{% endif %}'
   args:
     chdir: /home/frappe/agent
 

press/playbooks/roles/agent_monitoring_setup/tasks/main.yml

@@ -21,7 +21,7 @@
 - name: Generate Agent Configuration File
   become: yes
   become_user: frappe
-  command: '/home/frappe/agent/env/bin/agent setup config --name {{ server }} --workers {{ workers }} {% if proxy_ip is defined and proxy_ip is truthy %}--proxy-ip {{ proxy_ip }}{% endif %} {% if agent_sentry_dsn is defined and agent_sentry_dsn is truthy %}--sentry-dsn {{ agent_sentry_dsn }}{% endif %}'
+  command: '/home/frappe/agent/env/bin/agent setup config --name {{ server }} --workers {{ workers }} {% if proxy_ip is defined and proxy_ip is truthy %}--proxy-ip {{ proxy_ip }}{% endif %} {% if agent_sentry_dsn is defined and agent_sentry_dsn is truthy %}--sentry-dsn {{ agent_sentry_dsn }}{% endif %} {% if agent_token is defined and agent_token is truthy %}--agent-token {{ agent_token }}{% endif %}'
   args:
     chdir: /home/frappe/agent