Kong Rest Client

kong-rest-client.py

import requests


API_KEY = 'admin-consumer-api-key'
API_BASE = 'https://api.amzdata.io/admin'
CONSUMER_BASE = f'{API_BASE}/consumers'
API_ACL_GROUP_NAME = 'paid-users'


class KongAdmin:

    def __init__(self):
        self.headers = {
            'X-API-KEY': API_KEY
        }
        self.client = requests.session()
    
    def _make_request(self, url, method='POST', data=None):
        if method == 'POST':
            return self.client.post(url, data, headers=self.headers)
        if method == 'GET':
            return self.client.get(url, headers=self.headers)
        if method == 'DELETE':
            return self.client.delete(url, headers=self.headers)
    
    def basic_setup(self):
        """
        Setup basic settings of the API.
        """
    
    def get_consumers(self):
        """
        Get all consumers.
        """
        res = self._make_request(CONSUMER_BASE, method='GET')
        if res.status_code == 200:
            return res.json()
        else:
            return None

    def create_consumer(self, username):
        """
        Create a consumer in API gateway.
        """
        res = self._make_request(CONSUMER_BASE, data={'username': username})
        if res.status_code == 201:
            consumer_id = res.json().get('id')
            return consumer_id
        return False
    
    def delete_consumer(self, consumer_id):
        """
        Delete a consumer from API gateway.
        """
        self._make_request(f'{CONSUMER_BASE}/{consumer_id}', method='DELETE')
    
    def create_api_key(self, consumer_id):
        """
        Create an API key for a consumer.
        """
        res = self._make_request(f'{CONSUMER_BASE}/{consumer_id}/key-auth')
        if res.status_code == 201:
            return res.json().get('key')
        return False
    
    def reset_api_key(self, consumer_id):
        """
        Reset a consumer's all API keys.
        """
        res = self._make_request(f'{CONSUMER_BASE}/{consumer_id}/key-auth', method='GET')
        if res.status_code == 200:
            keys = res.json().get('data')
            for key in keys:
                keyId = key.get('id')
                self._make_request(f'{CONSUMER_BASE}/{consumer_id}/key-auth/{keyId}', method='DELETE')
            return self.create_api_key(consumer_id)
        
        return False
    
    def enable_ratelimiting(self, consumer_id, **kwargs):
        """
        Enable rate limiting on a consumer.
        """

        PLUGINS_BASE = f'{CONSUMER_BASE}/{consumer_id}/plugins'

        data = {
            'name': 'rate-limiting'
        }
        if kwargs.get('second'):
            data['config.second'] = kwargs.get('second')
        
        if kwargs.get('hour'):
            data['config.hour'] = kwargs.get('hour')
        
        if kwargs.get('minute'):
            data['config.minute'] = kwargs.get('minute')
        
        # Delete if rate limiting enabled
        res = self._make_request(PLUGINS_BASE, method='GET')
        for plugin in res.json().get('data'):
            if plugin.get('name') == 'rate-limiting':
                pluginId = plugin.get('id')
                self._make_request(f'{PLUGINS_BASE}/{pluginId}', method='DELETE')
        
        res = self._make_request(PLUGINS_BASE, data=data)
        if res.status_code == 201:
            return True
        return False
    
    def modify_api_access(self, consumer_id, status=False):
        ACL_BASE = f'{API_BASE}/consumers/{consumer_id}/acls'
        group_name = API_ACL_GROUP_NAME
        if status == True:
            res = self._make_request(ACL_BASE, data={'group': group_name})
            if res.status_code == 201:
                return True
        else:
            res = self._make_request(ACL_BASE, method='GET')
            if res.status_code == 200:
                acls = res.json().get('data')
                for acl in acls:
                    if acl.get('group') == group_name:
                        group_id = acl.get('id')
                        self._make_request(f'{ACL_BASE}/{group_id}', method='DELETE')
            return True
        
        return False