organization.py

import random
import string
import logging
from typing import Dict, List, Union
from datetime import datetime

from brian_dashboard_manager.grafana.utils.request import AdminRequest, \
    TokenRequest

logger = logging.getLogger(__name__)


def switch_active_organization(request: AdminRequest, org_id: int):
    assert org_id

    logger.debug(f'Switched {str(request)} active organization to #{org_id}')
    return request.post(f'api/user/using/{org_id}', {})


def get_organizations(request: AdminRequest) -> List:
    return request.get('api/orgs')


def create_organization(request: AdminRequest, name: str) -> Union[Dict, None]:
    assert name

    result = request.post('api/orgs', json={
        'name': name
    })

    if result.get('message', '').lower() == 'organization created':
        id = result.get('orgId')
        logger.info(f'Created organization `{name}` with ID #{id}')
        return {'id': id, 'name': name}
    else:
        return None


def delete_organization(request: AdminRequest, id: int) -> bool:

    result = request.delete(f'api/orgs/{id}')

    return result.get('message', '').lower() == 'organization deleted'


def create_api_token(request: AdminRequest, org_id: int, key_data=None):
    characters = string.ascii_uppercase + string.digits
    name = ''.join(random.choices(characters, k=16))
    data = {
        'name': name,
        'role': 'Admin',
        'secondsToLive': 3600  # 60 minutes
    }
    if key_data:
        data.update(key_data)

    switch_active_organization(request, org_id)
    result = request.post('api/auth/keys', json=data)
    token_id = result.get('id')

    logger.debug(f'Created API token #{token_id} for organization #{org_id}')

    return result


def delete_api_token(request: AdminRequest, org_id: int, token_id: int):
    assert token_id

    switch_active_organization(request, org_id)
    result = request.delete(f'api/auth/keys/{token_id}')
    logger.debug(f'Deleted API token #{token_id} for organization #{org_id}')
    return result


def delete_expired_api_tokens(request: AdminRequest, org_id: int) -> bool:
    assert org_id

    tokens = request.get('api/auth/keys', params={'includeExpired': True})

    now = datetime.now()

    def is_expired(token):
        date = datetime.strptime(token['expiration'], '%Y-%m-%dT%H:%M:%SZ')
        return date < now

    expired_tokens = [t for t in tokens if 'expiration' in t and is_expired(t)]

    for token in expired_tokens:
        delete_api_token(request, org_id, token['id'])
    return True


def set_home_dashboard(request: TokenRequest, dashboard_id: int):
    r = request.put('api/org/preferences', json={
        'homeDashboardId': dashboard_id
    })
    return r and r.get('message') == 'Preferences updated'