submit.py 2.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677
  1. """API Gateway handler: validate API key, create job, start Step Functions."""
  2. import json
  3. import uuid
  4. import time
  5. import os
  6. import boto3
  7. from shared import get_user, api_response, JOBS_TABLE, S3_BUCKET
  8. sfn = boto3.client('stepfunctions')
  9. ddb = boto3.resource('dynamodb')
  10. def handler(event, context):
  11. # Parse body
  12. try:
  13. body = json.loads(event.get('body', '{}'))
  14. except Exception:
  15. return api_response(400, {'error': 'Invalid JSON body'})
  16. api_key = body.get('api_key') or event.get('headers', {}).get('x-api-key', '')
  17. curl_command = body.get('curl')
  18. cookies = body.get('cookies', '')
  19. if not api_key:
  20. return api_response(401, {'error': 'Missing api_key'})
  21. if not curl_command:
  22. return api_response(400, {'error': 'Missing curl'})
  23. # Validate user
  24. user = get_user(api_key)
  25. if not user:
  26. return api_response(401, {'error': 'Invalid API key'})
  27. # Create job
  28. job_id = str(uuid.uuid4())[:8]
  29. now = int(time.time())
  30. ttl = now + 7 * 86400 # 7 days
  31. table = ddb.Table(JOBS_TABLE)
  32. table.put_item(Item={
  33. 'job_id': job_id,
  34. 'api_key': api_key,
  35. 'email': user.get('email', ''),
  36. 'status': 'SUBMITTED',
  37. 'created_at': now,
  38. 'ttl': ttl,
  39. })
  40. # Store curl + cookies in S3 (avoid DynamoDB size limits)
  41. s3 = boto3.client('s3')
  42. s3.put_object(
  43. Bucket=S3_BUCKET,
  44. Key=f"jobs/{job_id}/input.json",
  45. Body=json.dumps({
  46. 'curl': curl_command,
  47. 'cookies': cookies,
  48. 'language': body.get('language', 'auto'),
  49. }),
  50. ContentType='application/json',
  51. )
  52. # Start Step Functions
  53. state_machine_arn = os.environ.get('STATE_MACHINE', '')
  54. sfn.start_execution(
  55. stateMachineArn=state_machine_arn,
  56. name=f"job-{job_id}",
  57. input=json.dumps({
  58. 'job_id': job_id,
  59. 'email': user.get('email', ''),
  60. 'user_name': user.get('name', ''),
  61. }),
  62. )
  63. return api_response(200, {
  64. 'job_id': job_id,
  65. 'message': 'Job submitted, you will receive an email when complete.',
  66. })