MAINHACK

Viewing / Editing File: countries.py

File is not writable. Editing disabled.
from peewee import DoesNotExist

import im360.subsys.webshield as webshield
from defence360agent.rpc_tools import ValidationError, lookup
from defence360agent.rpc_tools.utils import (generate_warnings,
                                             run_in_executor_decorator)
from defence360agent.utils import Scope
from im360.api.ips import CountryAPI
from im360.model.country import Country, CountryList


class CountriesEndpoints(lookup.RootEndpoints):
    SCOPE = Scope.IM360

    @lookup.bind('blacklist', 'country', 'list')
    @run_in_executor_decorator
    def blacklist_country_list(self, limit=None, offset=None, **kwargs):
        return (CountryList.fetch_count(**kwargs),
                CountryList.fetch(limit=limit, offset=offset,
                                  by_list=CountryList.BLACK, **kwargs))

    @lookup.bind('blacklist', 'country', 'add')
    async def blacklist_country_add(self, items, comment=None):
        return await self._country_add(CountryList.BLACK, items, comment)

    @lookup.bind('blacklist', 'country', 'edit')
    async def blacklist_country_edit(self, items, comment):
        return await self._country_edit(CountryList.BLACK, items, comment)

    @lookup.bind('blacklist', 'country', 'delete')
    async def blacklist_country_delete(self, items):
        return await self._country_delete(CountryList.BLACK, items)

    @lookup.bind('whitelist', 'country', 'list')
    @run_in_executor_decorator
    def whitelist_country_list(self, limit=None, offset=None, **kwargs):
        return (CountryList.fetch_count(**kwargs),
                CountryList.fetch(limit=limit, offset=offset,
                                  by_list=CountryList.WHITE, **kwargs))

    @lookup.bind('whitelist', 'country', 'edit')
    async def whitelist_country_edit(self, items, comment):
        return await self._country_edit(CountryList.WHITE, items, comment)

    @lookup.bind('whitelist', 'country', 'add')
    async def whitelist_country_add(self, items, comment=None):
        return await self._country_add(CountryList.WHITE, items, comment)

    @lookup.bind('whitelist', 'country', 'delete')
    async def whitelist_country_delete(self, items):
        return await self._country_delete(CountryList.WHITE, items)

    @run_in_executor_decorator
    def _get_country(self, country_code):
        try:
            Country.get(code=country_code)
        except DoesNotExist:
            return False
        return True

    async def _check_countries_exists(self, country_list):
        warning_countries = []
        for country_code in country_list:
            if not (await self._get_country(country_code=country_code)):
                warning_countries.append(country_code)

        if warning_countries:
            raise ValidationError(
                'Country does not exist {}'.format(
                    ', '.join(warning_countries))
            )

    async def _country_edit(self, listname, items, comment):
        await self._check_countries_exists(items)
        affected, not_affected = await CountryAPI.edit(
            items, comment=comment)
        return generate_warnings(
            affected, not_affected,
            dest_listname=listname,
            all_list=items,
            success_warning="{}/{} country(s) were successfully edited",
            failure_warning="Noop: unable to edit {} in {}",
            in_another_list_warning="Country {} is already in {} list"
        )

    async def _country_add(self, listname, items, comment):
        """Add Country to appropriate ipset and CountryList"""
        await self._check_countries_exists(items)
        affected, not_affected = await CountryAPI.block(
            items, listname=listname, comment=comment)
        if listname == CountryList.BLACK and affected:
            # Update webshield conf file only for blacklisted country
            await webshield.update_country_blacklist_config()

        return generate_warnings(
            affected, not_affected,
            dest_listname=listname,
            all_list=items,
            success_warning="{}/{} ip(s) were successfully added",
            failure_warning="Noop: unable to add {} from {}",
            in_another_list_warning="Country {} is already in {} list"
        )

    async def _country_delete(self, listname, items):
        affected, not_affected = await CountryAPI.unblock(
            items, listname=listname)
        if listname == CountryList.BLACK and affected:
            # Update webshield conf file only for blacklisted country
            await webshield.update_country_blacklist_config()
        return generate_warnings(
            affected, not_affected,
            dest_listname=listname,
            all_list=items,
            success_warning="{}/{} ip(s) were successfully deleted",
            failure_warning="Noop: unable to delete {} from {}",
            in_another_list_warning="Country {} is already in {} list"
        )
Cancel / Back
[ MAINHACK ]

Viewing / Editing File: countries.py

Server Info

System Features

User Info
