Skip to content

Commit 92e9a89

Browse files
danfmaiadanfmaia
committed
feat: add API key purging functionality and improve key management
1 parent 4a654a6 commit 92e9a89

File tree

4 files changed

+304
-15
lines changed

4 files changed

+304
-15
lines changed

‎README.md

+39-1
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,44 @@ The CodeQuery API consists of two components:
122122
make logs # View container logs
123123
```
124124

125+
### API Key Management
126+
127+
The Gateway service provides functionality to manage your API keys. Here's how to perform common operations:
128+
129+
#### Generating a New API Key
130+
131+
```bash
132+
# Generate a key with default settings (30-day expiration, 60 requests/minute)
133+
curl -X POST https://codequery.dev/api-keys/generate
134+
135+
# Generate a key with custom settings
136+
curl -X POST -H "Content-Type: application/json" -d '{
137+
"expiration_days": 90,
138+
"requests_per_minute": 100
139+
}' https://codequery.dev/api-keys/generate
140+
```
141+
142+
#### Purging Your API Key
143+
144+
When you no longer need an API key, you can purge it from the server. This will:
145+
146+
- Delete the API key from the system
147+
- Remove all associated data
148+
- Invalidate any active ngrok tunnels for this key
149+
150+
To purge your API key:
151+
152+
```bash
153+
# Replace 'your-api-key' with your actual API key
154+
curl -X DELETE -H "X-API-KEY: your-api-key" https://codequery.dev/api-keys/your-api-key
155+
```
156+
157+
Note:
158+
159+
- You can only purge your own API key
160+
- Once purged, a key cannot be recovered
161+
- Make sure to generate a new key before purging if you plan to continue using the service
162+
125163
### Environment Variables
126164

127165
The `.env` file created by `make init` contains all necessary configuration. The main variables you'll need to set are:
@@ -168,7 +206,7 @@ For more detailed information about the API endpoints and advanced usage, see th
168206
3. **Domain Setup**: Consider using a custom domain for access.
169207
4. **SSL/TLS Configuration**: Use services like Let's Encrypt to secure the server.
170208

171-
## API Endpoints
209+
## Main API Endpoints
172210

173211
### 1. **Retrieve Project Structure**
174212

‎gateway/gateway.py

+101-13
Original file line numberDiff line numberDiff line change
@@ -86,30 +86,24 @@ async def api_key_validator(request: Request, call_next):
8686
"""
8787
Middleware to validate API keys and dynamically set the ngrok URL for each request.
8888
"""
89-
# Skip authentication for root and API key generation endpoints
90-
if request.url.path == "/" or request.url.path == "/api-keys/generate":
89+
# Skip authentication for root, API key generation, and purge endpoints
90+
if request.url.path == "/" or request.url.path == "/api-keys/generate" or request.url.path.startswith("/api-keys/") and request.method == "DELETE":
9191
return await call_next(request)
9292

9393
api_key = request.headers.get("x-api-key")
9494
if not api_key:
9595
return JSONResponse(status_code=401, content={"detail": "Missing API Key"})
9696

9797
try:
98-
# First check in-memory keys for faster validation
99-
if api_key not in self.api_keys:
100-
# If not in memory, try loading from S3
101-
api_keys = self.s3_manager.load_encrypted_api_keys() or {}
102-
if api_key not in api_keys:
103-
return JSONResponse(status_code=401, content={"detail": "Invalid API Key"})
104-
105-
# Add to in-memory cache if found in S3
106-
self.api_keys[api_key] = f"User{len(self.api_keys) + 1}"
107-
108-
# Load the latest key data from S3 for expiration and rate limit checks
98+
# First check in S3 for faster validation
10999
api_keys = self.s3_manager.load_encrypted_api_keys() or {}
110100
if api_key not in api_keys:
111101
return JSONResponse(status_code=401, content={"detail": "Invalid API Key"})
112102

103+
# Update in-memory cache if needed
104+
if api_key not in self.api_keys:
105+
self.api_keys[api_key] = f"User{len(self.api_keys) + 1}"
106+
113107
key_data = api_keys[api_key]
114108
current_time = datetime.datetime.utcnow()
115109

@@ -434,6 +428,100 @@ async def generate_api_key(request: Request):
434428
detail="Failed to generate API key. Please try again later."
435429
)
436430

431+
@self.app.delete("/api-keys/{api_key}")
432+
async def purge_api_key(api_key: str, request: Request):
433+
"""
434+
Purge all data associated with a specific API key.
435+
This includes removing the key from api_keys.json and ngrok_urls.json.
436+
Users can purge their own keys, while admin can purge any key.
437+
"""
438+
try:
439+
# URL decode the API key
440+
api_key = unquote_plus(api_key)
441+
442+
# Get the request API key
443+
request_api_key = request.headers.get("x-api-key")
444+
if not request_api_key:
445+
raise HTTPException(
446+
status_code=401,
447+
detail="Missing API Key"
448+
)
449+
450+
# Get the admin key from environment variables
451+
admin_key = os.getenv("ADMIN_API_KEY")
452+
453+
# Load current API keys
454+
api_keys = self.s3_manager.load_encrypted_api_keys() or {}
455+
456+
# Check if the key exists
457+
if api_key not in api_keys:
458+
raise HTTPException(
459+
status_code=404,
460+
detail=f"API key {api_key} not found"
461+
)
462+
463+
# Check authorization:
464+
# 1. Admin can purge any key except admin key
465+
# 2. Users can only purge their own key
466+
if request_api_key != admin_key and request_api_key != api_key:
467+
raise HTTPException(
468+
status_code=401,
469+
detail="Unauthorized. You can only purge your own API key."
470+
)
471+
472+
# Check if trying to delete admin key
473+
if api_key == admin_key:
474+
raise HTTPException(
475+
status_code=403,
476+
detail="Cannot purge admin API key"
477+
)
478+
479+
# Store key data for audit log
480+
purged_key_data = api_keys[api_key]
481+
482+
# Remove the key from api_keys.json
483+
del api_keys[api_key]
484+
self.s3_manager.store_encrypted_api_keys(api_keys)
485+
486+
# Remove the key from ngrok_urls.json
487+
try:
488+
self.s3_manager.update_ngrok_url(api_key, None)
489+
except Exception as e:
490+
self.logger.error(
491+
f"Error removing ngrok URL for {api_key}: {str(e)}")
492+
# Continue with the purge even if ngrok URL removal fails
493+
494+
# Invalidate the in-memory cache
495+
self.invalidate_ngrok_cache(api_key)
496+
497+
# Log the purge operation for audit trail
498+
self.logger.info(
499+
"API key purged: %s, Created: %s, Last Used: %s, Total Requests: %d",
500+
api_key,
501+
purged_key_data.get("created_at"),
502+
purged_key_data.get("last_used"),
503+
purged_key_data.get("total_requests", 0)
504+
)
505+
506+
return {
507+
"status": "success",
508+
"message": f"API key {api_key} purged successfully",
509+
"purged_data": {
510+
"created_at": purged_key_data.get("created_at"),
511+
"last_used": purged_key_data.get("last_used"),
512+
"total_requests": purged_key_data.get("total_requests", 0)
513+
}
514+
}
515+
516+
except HTTPException:
517+
raise
518+
except Exception as e:
519+
self.logger.error(f"Error purging API key: {str(e)}")
520+
raise HTTPException(
521+
status_code=500,
522+
detail=f"Failed to purge API key: {str(e)}"
523+
) from e
524+
437525

438526
# Create an instance of the GatewayAPI class
439527
gateway_instance = GatewayAPI()

‎gateway/template.env

+2-1
Original file line numberDiff line numberDiff line change
@@ -10,4 +10,5 @@ EC2_HOST="" # Your EC2 instance's public DNS (e.g., ec2-xx-xx-xxx-xxx.compute.a
1010
KEY_PATH="./secrets/codequery-keypair.pem" # Path to your SSH private key for EC2 access
1111

1212
# API Configuration
13-
API_KEYS= # Comma-separated list of initial API keys (optional)
13+
API_KEYS= # Comma-separated list of initial API keys (optional)
14+
ADMIN_API_KEY= # Admin API key for managing other API keys

‎gateway/tests/test_gateway.py

+162
Original file line numberDiff line numberDiff line change
@@ -642,3 +642,165 @@ def test_ngrok_url_update_error(self):
642642
self.assertEqual(response.status_code, 500)
643643
self.assertEqual(response.json(), {
644644
"detail": "Error updating ngrok URL"})
645+
646+
@patch.dict('os.environ', {'ADMIN_API_KEY': 'admin-key'})
647+
def test_purge_api_key(self):
648+
"""Test API key purge endpoint."""
649+
# Set up test data
650+
test_key = "test-key-to-purge"
651+
test_key_data = {
652+
"created_at": "2024-02-14T10:00:00",
653+
"last_used": "2024-02-14T11:00:00",
654+
"expires_at": None,
655+
"rate_limit": {
656+
"requests_per_minute": 60,
657+
"current_minute": None,
658+
"minute_requests": 0
659+
},
660+
"total_requests": 10
661+
}
662+
663+
# Set up S3 manager mock
664+
patcher = patch.object(self.gateway_instance, 's3_manager')
665+
mock_s3_manager = patcher.start()
666+
mock_s3_manager.load_encrypted_api_keys.return_value = {
667+
test_key: test_key_data,
668+
"admin-key": {
669+
"created_at": "2024-02-14T10:00:00",
670+
"last_used": None,
671+
"expires_at": None,
672+
"rate_limit": {
673+
"requests_per_minute": 60,
674+
"current_minute": None,
675+
"minute_requests": 0
676+
},
677+
"total_requests": 0
678+
}
679+
}
680+
self.addCleanup(patcher.stop)
681+
682+
# Test successful self-purge by user
683+
headers = {"x-api-key": test_key}
684+
response = self.client.delete(f"/api-keys/{test_key}", headers=headers)
685+
self.assertEqual(response.status_code, 200)
686+
response_data = response.json()
687+
self.assertEqual(response_data["status"], "success")
688+
self.assertEqual(response_data["purged_data"]["total_requests"], 10)
689+
690+
# Verify S3 manager calls for self-purge
691+
mock_s3_manager.store_encrypted_api_keys.assert_called()
692+
mock_s3_manager.update_ngrok_url.assert_called_with(test_key, None)
693+
694+
# Reset call counts and mock data
695+
mock_s3_manager.store_encrypted_api_keys.reset_mock()
696+
mock_s3_manager.update_ngrok_url.reset_mock()
697+
mock_s3_manager.load_encrypted_api_keys.return_value = {
698+
test_key: test_key_data,
699+
"admin-key": {
700+
"created_at": "2024-02-14T10:00:00",
701+
"last_used": None,
702+
"expires_at": None,
703+
"rate_limit": {
704+
"requests_per_minute": 60,
705+
"current_minute": None,
706+
"minute_requests": 0
707+
},
708+
"total_requests": 0
709+
}
710+
}
711+
712+
# Test successful purge by admin
713+
headers = {"x-api-key": "admin-key"}
714+
response = self.client.delete(f"/api-keys/{test_key}", headers=headers)
715+
self.assertEqual(response.status_code, 200)
716+
response_data = response.json()
717+
self.assertEqual(response_data["status"], "success")
718+
self.assertEqual(response_data["purged_data"]["total_requests"], 10)
719+
720+
# Verify S3 manager calls
721+
mock_s3_manager.store_encrypted_api_keys.assert_called()
722+
mock_s3_manager.update_ngrok_url.assert_called_with(test_key, None)
723+
724+
@patch.dict('os.environ', {'ADMIN_API_KEY': 'admin-key'})
725+
def test_purge_api_key_unauthorized(self):
726+
"""Test API key purge endpoint with unauthorized access."""
727+
# Set up test data
728+
test_key = "test-key-to-purge"
729+
other_key = "other-key"
730+
test_key_data = {
731+
"created_at": "2024-02-14T10:00:00",
732+
"last_used": None,
733+
"expires_at": None,
734+
"rate_limit": {
735+
"requests_per_minute": 60,
736+
"current_minute": None,
737+
"minute_requests": 0
738+
},
739+
"total_requests": 0
740+
}
741+
742+
# Set up S3 manager mock
743+
patcher = patch.object(self.gateway_instance, 's3_manager')
744+
mock_s3_manager = patcher.start()
745+
mock_s3_manager.load_encrypted_api_keys.return_value = {
746+
test_key: test_key_data,
747+
other_key: test_key_data
748+
}
749+
self.addCleanup(patcher.stop)
750+
751+
# Test with another user's key (not admin, not self)
752+
headers = {"x-api-key": other_key}
753+
response = self.client.delete(f"/api-keys/{test_key}", headers=headers)
754+
self.assertEqual(response.status_code, 401)
755+
self.assertEqual(
756+
response.json()["detail"], "Unauthorized. You can only purge your own API key.")
757+
758+
# Test without API key
759+
response = self.client.delete(f"/api-keys/{test_key}")
760+
self.assertEqual(response.status_code, 401)
761+
self.assertEqual(response.json()["detail"], "Missing API Key")
762+
763+
@patch.dict('os.environ', {'ADMIN_API_KEY': 'admin-key'})
764+
def test_purge_admin_key(self):
765+
"""Test attempt to purge admin API key."""
766+
# Set up S3 manager mock
767+
patcher = patch.object(self.gateway_instance, 's3_manager')
768+
mock_s3_manager = patcher.start()
769+
mock_s3_manager.load_encrypted_api_keys.return_value = {
770+
"admin-key": {
771+
"created_at": "2024-02-14T10:00:00",
772+
"last_used": None,
773+
"expires_at": None,
774+
"rate_limit": {
775+
"requests_per_minute": 60,
776+
"current_minute": None,
777+
"minute_requests": 0
778+
},
779+
"total_requests": 0
780+
}
781+
}
782+
self.addCleanup(patcher.stop)
783+
784+
# Test attempt to purge admin key
785+
headers = {"x-api-key": "admin-key"}
786+
response = self.client.delete("/api-keys/admin-key", headers=headers)
787+
self.assertEqual(response.status_code, 403)
788+
self.assertEqual(
789+
response.json()["detail"], "Cannot purge admin API key")
790+
791+
@patch.dict('os.environ', {'ADMIN_API_KEY': 'admin-key'})
792+
def test_purge_nonexistent_key(self):
793+
"""Test attempt to purge a non-existent API key."""
794+
# Set up S3 manager mock
795+
patcher = patch.object(self.gateway_instance, 's3_manager')
796+
mock_s3_manager = patcher.start()
797+
mock_s3_manager.load_encrypted_api_keys.return_value = {}
798+
self.addCleanup(patcher.stop)
799+
800+
# Test purge of non-existent key
801+
headers = {"x-api-key": "admin-key"}
802+
response = self.client.delete(
803+
"/api-keys/nonexistent-key", headers=headers)
804+
self.assertEqual(response.status_code, 404)
805+
self.assertEqual(
806+
response.json()["detail"], "API key nonexistent-key not found")

0 commit comments

Comments
 (0)