It all works.
modified: app.py modified: templates/display_agent_jwt.html.j2 modified: templates/manage_namespace.html.j2 modified: test_agent.sh
This commit is contained in:
parent
a8b0ee313d
commit
5af4d5a7e4
4 changed files with 74 additions and 45 deletions
111
app.py
111
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"<Agent(name='{self.name}', namespace_id='{self.namespace_id}')>"
|
||||
return f"<Agent(name='{self.name}', namespace_id='{self.namespace_id}', status='{self.status}')>"
|
||||
|
||||
|
||||
################################################################################
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -4,9 +4,13 @@
|
|||
|
||||
{% block content %}
|
||||
<h1>JWT for Agent '{{ agent_name }}'</h1>
|
||||
{% if message %}
|
||||
<p>{{ message }}</p>
|
||||
{% endif %}
|
||||
<p>Please store this JWT securely. It will not be shown again.</p>
|
||||
<pre>{{ jwt_token }}</pre>
|
||||
<p>You can use this JWT to authenticate your agent when sending logs to the namespace '{{ namespace.name }}'.</p>
|
||||
|
||||
<p><a href="{{ request.route_url('manage_namespace', namespace_short_id=namespace.short_id) }}">Back to Manage Namespace</a></p>
|
||||
{% endblock %}
|
||||
|
||||
|
|
|
|||
|
|
@ -44,7 +44,7 @@
|
|||
<ul>
|
||||
{% for agent in agents %}
|
||||
<li>
|
||||
{{ agent.name }}
|
||||
{{ agent.name }} ({{ agent.status|capitalize }})
|
||||
<form method="POST" action="{{ request.route_url('revoke_agent', namespace_short_id=namespace.short_id) }}" style="display:inline;">
|
||||
<input type="hidden" name="agent_id" value="{{ agent.id }}">
|
||||
<button type="submit" onclick="return confirm('Revoke access for {{ agent.name }}?')">Revoke Access</button>
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue