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 @@