From 5af4d5a7e4631604c07bd2c5499cc00cfa753e75 Mon Sep 17 00:00:00 2001 From: Russell Ballestrini Date: Sat, 11 Jan 2025 20:26:01 -0500 Subject: [PATCH] It all works. modified: app.py modified: templates/display_agent_jwt.html.j2 modified: templates/manage_namespace.html.j2 modified: test_agent.sh --- app.py | 111 +++++++++++++++++----------- templates/display_agent_jwt.html.j2 | 4 + templates/manage_namespace.html.j2 | 2 +- test_agent.sh | 2 +- 4 files changed, 74 insertions(+), 45 deletions(-) diff --git a/app.py b/app.py index 9b3ddb4..08e5adf 100644 --- a/app.py +++ b/app.py @@ -1,6 +1,8 @@ ############################################################################### -# app.py - PyraLogs Application with Agent-Specific JWT Authentication and Revocation -# Users do not have token_version as only agents use JWTs +# app.py - PyraLogs Application with Enhanced Agent Management +# - Prevents duplicate agent names within a namespace +# - Handles agent revocation by updating agent status +# - Provides JWT regeneration for existing agents ############################################################################### import os import base64 @@ -262,11 +264,12 @@ class Agent(Base): namespace_id = Column(String, ForeignKey("namespaces.id")) token_version = Column(Integer, default=0) created_at = Column(DateTime, default=datetime.datetime.utcnow) + status = Column(String, default="active") # 'active' or 'revoked' namespace = relationship("Namespace", back_populates="agents") def __repr__(self): - return f"" + return f"" ################################################################################ @@ -565,6 +568,7 @@ def logout_view(request): # JWT Helper Functions ################################################################################ + def generate_jwt_token(agent): """Generate a JWT for the given agent without an expiration time.""" payload = { @@ -744,8 +748,12 @@ def manage_namespace_view(request): } ) - # Get agents associated with the namespace - agents = s.query(Agent).filter(Agent.namespace_id == namespace.id).all() + # Get agents associated with the namespace (only active agents) + agents = ( + s.query(Agent) + .filter(Agent.namespace_id == namespace.id, Agent.status == "active") + .all() + ) return { "request": request, @@ -844,29 +852,50 @@ def generate_agent_jwt_view(request): if not agent_name: return Response("Agent name is required.", status=400) - # Generate UUID for the agent - agent_uuid = uuid.uuid4() - agent_id = str(agent_uuid) - - # Create an agent entry - agent = Agent( - id=agent_id, - name=agent_name, - namespace_id=namespace.id, - token_version=0, + # Check if an agent with the same name exists in the namespace + agent = ( + s.query(Agent) + .filter( + Agent.name == agent_name, + Agent.namespace_id == namespace.id, + ) + .first() ) - s.add(agent) - s.flush() + if agent: + # Agent exists, increment token_version and regenerate JWT + agent.token_version += 1 + agent.status = "active" # Ensure the agent is active + s.flush() + jwt_token = generate_jwt_token(agent) + message = f"A new JWT has been generated for existing agent '{agent.name}'. Any previous tokens have been revoked." + else: + # Generate UUID for the new agent + agent_uuid = uuid.uuid4() + agent_id = str(agent_uuid) - # Generate JWT for the agent including the namespace_id - jwt_token = generate_jwt_token(agent) + # Create a new agent entry + agent = Agent( + id=agent_id, + name=agent_name, + namespace_id=namespace.id, + token_version=0, + status="active", + ) + + s.add(agent) + s.flush() + + # Generate JWT for the agent including the namespace_id + jwt_token = generate_jwt_token(agent) + message = f"A new agent '{agent.name}' has been created." return { "request": request, "namespace": namespace, "agent_name": agent_name, "jwt_token": jwt_token, + "message": message, } @@ -881,18 +910,22 @@ def revoke_agent_view(request): return Response("Agent ID is required.", status=400) # Get the agent - agent = s.query(Agent).filter( - Agent.id == agent_id, - Agent.namespace_id == namespace.id, - ).first() + agent = ( + s.query(Agent) + .filter( + Agent.id == agent_id, + Agent.namespace_id == namespace.id, + ) + .first() + ) if not agent: return Response("Agent not found.", status=404) - # Increment the agent's token_version to revoke its token - agent.token_version += 1 + # Set the agent's status to 'revoked' to hide it from the dashboard + agent.status = "revoked" s.flush() - request.session.flash(f"Access revoked for agent '{agent.name}'.") + request.session.flash(f"Agent '{agent.name}' has been revoked.") return HTTPFound( location=request.route_url( @@ -927,6 +960,10 @@ def log_webhook_view(request): if not agent: return Response("Agent not found.", status=404) + # Check agent status + if agent.status != "active": + return Response("Agent is revoked or inactive.", status=401) + # Check token_version if payload.get("token_version") != agent.token_version: return Response("Token has been revoked.", status=401) @@ -967,19 +1004,13 @@ def log_webhook_view(request): return Response("Log entry created.", status=201) -@view_config( - route_name="view_namespace_logs", renderer="view_namespace_logs.html.j2" -) +@view_config(route_name="view_namespace_logs", renderer="view_namespace_logs.html.j2") @reader_required def view_namespace_logs_view(request): namespace = request.namespace namespace_dbsession = request.namespace_dbsession - logs = ( - namespace_dbsession.query(LogEntry) - .order_by(LogEntry.timestamp.desc()) - .all() - ) + logs = namespace_dbsession.query(LogEntry).order_by(LogEntry.timestamp.desc()).all() return { "request": request, @@ -1064,14 +1095,10 @@ def main(global_config=None, **settings): # Namespace Management config.add_route("create_namespace", "/namespace/create") - config.add_route( - "manage_namespace", "/namespace/{namespace_short_id}/manage" - ) + config.add_route("manage_namespace", "/namespace/{namespace_short_id}/manage") config.add_route("update_namespace", "/namespace/{namespace_short_id}/update") config.add_route("invite_user", "/namespace/{namespace_short_id}/invite") - config.add_route( - "revoke_agent", "/namespace/{namespace_short_id}/revoke_agent" - ) + config.add_route("revoke_agent", "/namespace/{namespace_short_id}/revoke_agent") # Route to generate JWTs for agents config.add_route( "generate_agent_jwt", "/namespace/{namespace_short_id}/generate_agent_jwt" @@ -1079,9 +1106,7 @@ def main(global_config=None, **settings): # Logging config.add_route("log_webhook", "/webhook") - config.add_route( - "view_namespace_logs", "/namespace/{namespace_short_id}/logs" - ) + config.add_route("view_namespace_logs", "/namespace/{namespace_short_id}/logs") config.scan() return config.make_wsgi_app() diff --git a/templates/display_agent_jwt.html.j2 b/templates/display_agent_jwt.html.j2 index 4b73081..fb12467 100644 --- a/templates/display_agent_jwt.html.j2 +++ b/templates/display_agent_jwt.html.j2 @@ -4,9 +4,13 @@ {% block content %}

JWT for Agent '{{ agent_name }}'

+{% if message %} +

{{ message }}

+{% endif %}

Please store this JWT securely. It will not be shown again.

{{ jwt_token }}

You can use this JWT to authenticate your agent when sending logs to the namespace '{{ namespace.name }}'.

Back to Manage Namespace

{% endblock %} + diff --git a/templates/manage_namespace.html.j2 b/templates/manage_namespace.html.j2 index f8fbbaf..a8691b8 100644 --- a/templates/manage_namespace.html.j2 +++ b/templates/manage_namespace.html.j2 @@ -44,7 +44,7 @@
    {% for agent in agents %}
  • - {{ agent.name }} + {{ agent.name }} ({{ agent.status|capitalize }})
    diff --git a/test_agent.sh b/test_agent.sh index 5447d46..361cbc2 100644 --- a/test_agent.sh +++ b/test_agent.sh @@ -9,7 +9,7 @@ # Variables BASE_URL="http://127.0.0.1:6544" -JWT_TOKEN="eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJhZ2VudF9pZCI6IjMzMmM3ZjY2LWMzNTktNDg4Yi05MTIxLWIwN2Y1N2EyMTA3MCIsImFnZW50X25hbWUiOiJ0ZXN0LTkwMDAiLCJuYW1lc3BhY2VfaWQiOiI3YmM3Mjk5Yy1hYTgxLTRhYjQtYmNmMy00OTU5NzNhZjk3NzEiLCJ0b2tlbl92ZXJzaW9uIjowLCJpYXQiOjE3MzY2NDQyNjB9.fTZuLTFARjgTXIUtg2zGkfcZp3sAJYAasSa6MFxK-Gk" +JWT_TOKEN="eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJhZ2VudF9pZCI6IjcxZmQ0ZjA0LTQ2NzgtNGEyNi05ZmU4LTkxNjU2OTNmOTBmOSIsImFnZW50X25hbWUiOiJ0ZXN0LTkwMDAiLCJuYW1lc3BhY2VfaWQiOiJhZGQwYWQ1YS0zNjE0LTRhODgtOTRmZC0xNGQ4ZDBkYjIxYjUiLCJ0b2tlbl92ZXJzaW9uIjoyLCJpYXQiOjE3MzY2NDUxMDF9.WPsitdEYcJ6gpvExO8BmddJdmqBlczqQAD1SzHB0Y1M" # Log message from the first argument LOG_MESSAGE="$1"