feat: Enhance press-agent communication

.cspell.json

@@ -483,6 +483,7 @@
         "phpmyadmin",
         "pids",
         "Pjpw",
+        "pkcs",
         "pkgs",
         "pmadb",
         "pmezard",
@@ -593,6 +594,7 @@
         "sprintf",
         "squashfs",
         "Srednekolymsk",
+        "srem",
         "Starke",
         "stdc",
         "stime",

press/agent.py

@@ -3,16 +3,19 @@
 from __future__ import annotations
 
 import _io  # type: ignore
+import base64
 import json
 import os
 import re
+import time
 from contextlib import suppress
 from datetime import date
 from typing import TYPE_CHECKING, Any, Literal
 
 import frappe
 import frappe.utils
 import requests
+from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
 from frappe.utils.password import get_decrypted_password
 from requests.exceptions import HTTPError
 
@@ -26,6 +29,7 @@
 if TYPE_CHECKING:
 	from io import BufferedReader
 
+	from press.press.doctype.agent_auth.agent_auth import AgentAuth
 	from press.press.doctype.agent_job.agent_job import AgentJob
 	from press.press.doctype.app_patch.app_patch import AgentPatchConfig, AppPatch
 	from press.press.doctype.bench.bench import Bench
@@ -949,13 +953,105 @@ 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_agent_public_key(self):
+		key = f"{self.server}_agent_public_key"
+
+		public_key = frappe.cache().get_value(key)
+
+		if not public_key:
+			agent_auth: AgentAuth = frappe.get_doc(
+				"Agent Auth",
+				self.server,
+			)
+
+			public_key = agent_auth.public_key
+
+			if not frappe.cache().get_value(f"{self.server}_regenerate_public_key"):
+				# Don't set cache while regenerating. Old public key may get cached again.
+				frappe.cache().set_value(
+					key,
+					public_key,
+					expires_in_sec=3600,
+				)
+
+		return public_key
+
+	def get_regenerate_public_key(self):
+		key = f"{self.server}_regenerate_public_key"
+
+		regenerate_key = frappe.cache().get_value(key)
+		if not regenerate_key:
+			agent_auth: AgentAuth = frappe.get_doc(
+				"Agent Auth",
+				self.server,
+			)
+
+			if agent_auth.regenerate_public_key:
+				agent_auth.regenerate_public_key = None
+				agent_auth.save(ignore_permissions=True)
+
+			return None
+
+		return regenerate_key
+
+	def _verify_request_token(self, token: str):
+		from cryptography.exceptions import InvalidSignature
+
+		payload_b64, signature_b64 = token.split(".")
+
+		payload_bytes = base64.urlsafe_b64decode(payload_b64 + "==")
+		signature = base64.urlsafe_b64decode(signature_b64 + "==")
+
+		public_keys = [self.get_agent_public_key()]
+
+		regenerate_key = self.get_regenerate_public_key()
+
+		if regenerate_key:
+			public_keys.append(regenerate_key)
+
+		verified = False
+
+		for key in public_keys:
+			try:
+				public_key = Ed25519PublicKey.from_public_bytes(base64.b64decode(key))
+
+				public_key.verify(signature, payload_bytes)
+
+				verified = True
+				break
+
+			except InvalidSignature:
+				pass
+
+		if not verified:
+			raise ValueError("Invalid token signature")
+
+		payload = json.loads(payload_bytes)
+
+		if payload["exp"] < (time.time() - 60):
+			raise ValueError("Token expired")
+
+		if payload["server"] != self.server:
+			raise ValueError("Invalid server")
+
+		return True
+
+	def extract_and_verify_token(self, token):
+		if not token:
+			raise ValueError("Unsigned request from agent")
+
+		self._verify_request_token(
+			token=token,
+		)
+
 	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 +1122,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

press/api/agent_auth.py

@@ -0,0 +1,13 @@
+import frappe
+
+from press.agent import Agent
+
+
+def verify_agent(server: str):
+	agent_token = frappe.request.headers.get("X-Agent-Token")
+
+	if not agent_token:
+		frappe.throw_permission_error()
+
+	agent = Agent(server)
+	agent.extract_and_verify_token(agent_token)

press/api/callbacks.py

@@ -3,11 +3,13 @@
 from __future__ import annotations
 
 import ipaddress
+import json
 
 import frappe
 from frappe.rate_limiter import rate_limit
 
 from press.agent import Agent
+from press.api.agent_auth import verify_agent
 from press.press.doctype.agent_job.agent_job import handle_polled_job
 from press.utils import log_error
 
@@ -113,8 +115,40 @@ def callback(job_id: str | None = None):
 	if not server:
 		frappe.throw("Not permitted", frappe.ValidationError)
 
+	verify_agent(server)
+
 	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=10, seconds=60)
+def update_job(job, server):
+	flag = frappe.db.get_single_value("Press Settings", "push_feature")
+	if not flag:
+		return
+
+	if not job:
+		return
+
+	verify_agent(server)
+
+	job = json.loads(job)
+
+	job_doc = frappe.get_value(
+		"Agent Job",
+		fieldname=[
+			"name",
+			"job_id",
+			"status",
+			"callback_failure_count",
+			"job_type",
+		],
+		filters={"job_id": job["id"]},
+		as_dict=True,
+	)
+
+	handle_polled_job(polled_job=job, job=job_doc)

press/api/monitoring.py

@@ -7,6 +7,7 @@
 import frappe
 from frappe.rate_limiter import rate_limit
 
+from press.api.agent_auth import verify_agent
 from press.exceptions import AlertRuleNotEnabled
 from press.press.doctype.monitor_server.monitor_server import get_monitor_server_ips
 from press.utils import log_error, servers_using_alternative_port_for_communication
@@ -159,7 +160,9 @@ def get_targets_method_rate_limit() -> int:
 
 @frappe.whitelist(allow_guest=True)
 @rate_limit(limit=get_targets_method_rate_limit, seconds=MONITORING_ENDPOINT_RATE_LIMIT_WINDOW_SECONDS)
-def targets(token: str | None = None):
+def targets(server: str, token: str | None = None):
+	verify_agent(server)
+
 	if not token:
 		frappe.throw_permission_error()
 	monitor_token = frappe.db.get_single_value("Press Settings", "monitor_token", cache=True)

press/api/server.py

@@ -13,6 +13,7 @@
 from frappe.utils.caching import redis_cache
 from frappe.utils.password import get_decrypted_password
 
+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
@@ -974,20 +975,16 @@ def rename(name, title):
 
 
 @frappe.whitelist(allow_guest=True)
-def benches_are_idle(server: str, access_token: str) -> None:
+def benches_are_idle(server: str) -> 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
+	verify_agent(server)
 
 	primary_server, is_server_scaled_up = frappe.db.get_value(
 		"Server", {"secondary_server": server}, ["name", "scaled_up"]

press/hooks.py

@@ -197,6 +197,7 @@
 scheduler_events = {
 	"weekly_long": ["press.press.doctype.marketplace_app.events.auto_review_for_missing_steps"],
 	"daily": [
+		"press.press.doctype.agent_auth.agent_auth.regenerate_token",
 		"press.experimental.doctype.referral_bonus.referral_bonus.credit_referral_bonuses",
 		"press.press.doctype.log_counter.log_counter.record_counts",
 		"press.press.doctype.incident.incident.notify_ignored_servers",
@@ -354,6 +355,7 @@
 			"press.workflow_engine.doctype.press_workflow.press_workflow.retry_workflow_callbacks",
 		],
 		"* * * * *": [
+			"press.press.doctype.agent_job.agent_job.retry_poll",
 			"press.press.doctype.virtual_disk_snapshot.virtual_disk_snapshot.sync_physical_backup_snapshots",
 			"press.workflow_engine.doctype.press_workflow_task.press_workflow_task.retry_tasks",
 			"press.press.doctype.deploy_candidate_build.deploy_candidate_build.run_scheduled_builds",

press/playbooks/database.yml

@@ -13,6 +13,7 @@
     - role: mariadb_memory_allocator
     - role: nginx
     - role: agent
+    - role: setup_agent_auth
     - role: node_exporter
     - role: mysqld_exporter
     - role: deadlock_logger

press/playbooks/proxy.yml

@@ -9,6 +9,7 @@
     - role: user
     - role: nginx
     - role: agent
+    - role: setup_agent_auth
     - role: proxy
     - role: node_exporter
     - role: user_ssh_certificate

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

@@ -0,0 +1,6 @@
+---
+- name: Write agent token
+  copy:
+    content: "{{ agent_token }}"
+    dest: "/home/frappe/agent/agent.token"
+    mode: '0600'

press/playbooks/self_hosted.yml

@@ -9,6 +9,7 @@
     - role: user
     - role: nginx
     - role: agent
+    - role: setup_agent_auth
     - role: bench
     - role: docker
     - role: node_exporter

press/playbooks/self_hosted_db.yml

@@ -10,6 +10,7 @@
     - role: mariadb
     - role: nginx
     - role: agent
+    - role: setup_agent_auth
     - role: node_exporter
     - role: mysqld_exporter
     - role: deadlock_logger

press/playbooks/self_hosted_proxy.yml

@@ -9,5 +9,6 @@
     - role: user
     - role: nginx
     - role: agent
+    - role: setup_agent_auth
     - role: proxy
     - role: docker

press/playbooks/server.yml

@@ -10,6 +10,7 @@
     - role: user
     - role: nginx
     - role: agent
+    - role: setup_agent_auth
     - role: mount
     - role: bench
     - role: docker

press/playbooks/setup_agent_auth.yml

@@ -0,0 +1,8 @@
+---
+- name: Setup agent auth
+  hosts: all
+  become: yes
+  become_user: frappe
+  gather_facts: no
+  roles:
+    - role: setup_agent_auth

press/press/doctype/agent_auth/agent_auth.js

@@ -0,0 +1,8 @@
+// Copyright (c) 2026, Frappe and contributors
+// For license information, please see license.txt
+
+// frappe.ui.form.on("Agent Auth", {
+// 	refresh(frm) {
+
+// 	},
+// });