receive.py 1.85 KB
Newer Older
1 2 3 4 5 6
from flask_restful import Resource, request, abort, fields
from marshmallow import Schema, fields
import pika
import json


7
class ReceiveQuerySchema(Schema):
8 9 10 11 12 13 14 15 16 17
    client_id = fields.Str(required=True)
    secret = fields.Str(required=True)
    max_messages = fields.Int(required=False)


class Receive(Resource):
    clients = None
    schema = None

    def __init__(self):
18
        self.schema = ReceiveQuerySchema()
19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57
        with open("clients.json", "r") as clients_file:
            self.clients = json.load(clients_file)

    def get(self):
        errors = self.schema.validate(request.args)
        if errors:
            abort(400, str(errors))

        messages = []
        allow = False
        max_messages = request.args.get("max_messages", 10)
        client_id = request.args.get("client_id")
        inbox_queue = client_id + "-inbox"
        if client_id in self.clients:
            client = self.clients.get(client_id)
            if request.args.get("secret") == client.get("secret"):
                allow = True

        if allow:
            connection = pika.BlockingConnection(
                pika.ConnectionParameters(host="localhost")
            )
            channel = connection.channel()
            channel.queue_declare(queue=inbox_queue, durable=True)
            while len(messages) < max_messages:
                method_frame, header_frame, body = channel.basic_get(inbox_queue)
                if method_frame:
                    print(method_frame, header_frame, body)
                    channel.basic_ack(method_frame.delivery_tag)
                    messages.append(json.loads(body.decode()))
                else:
                    print("No message returned")
                    break
            channel.close()
            connection.close()

        else:
            abort(403, message="Invalid client credentials")
        return messages