Current File : //proc/thread-self/root/opt/imunify360/venv/lib64/python3.11/site-packages/im360/model/firewall.py
"""DB tables related to firewall functionality."""
import ipaddress
import itertools
import logging
import time
from datetime import timedelta
from enum import Enum
from functools import reduce
from ipaddress import IPv4Network, IPv6Network
from itertools import starmap
from operator import ior, itemgetter
from typing import (
Any,
Dict,
Iterable,
Iterator,
List,
Optional,
Sequence,
Tuple,
Type,
Union,
)
from blinker import Signal
from peewee import (
JOIN,
SQL,
BooleanField,
Case,
CharField,
Check,
CompositeKey,
DoesNotExist,
FloatField,
ForeignKeyField,
IntegerField,
PrimaryKeyField,
TextField,
fn,
prefetch,
)
from playhouse.shortcuts import model_to_dict
from defence360agent.contracts.messages import Reject
from defence360agent.model import Model, instance
from defence360agent.model.simplification import ApplyOrderBy
from defence360agent.utils import CHUNK_SIZE_SQL_QUERY, split_for_chunk, timeit
from im360.model.country import Country
from im360.utils.net import (
ALL,
TCP,
UDP,
IPNetwork,
pack_ip_network,
unpack_ip_network,
is_net,
)
from im360.utils.validate import IP, IPVersion, NumericIPVersion
IPListID = int # TypeAlias is 3.10+
logger = logging.getLogger(__name__)
IPV4_HOST_MASK = pack_ip_network(IPv4Network("0.0.0.0/32"))[1]
# minimal size of v6 network stored in db /64
IPV6_HOST_MASK = pack_ip_network(IPv6Network("::/64"))[1]
IPv4 = NumericIPVersion[IP.V4].value
IPv6 = NumericIPVersion[IP.V6].value
def _filter_ip_net_subnets(
model: Type[Model], packed_ip_net: Tuple[int, int, int]
):
"""
Filters ip addresses/networks contained in ip network net.
:param model: model to apply query
:param packed_ip_net: tuple of integers
:return: peewee expression
"""
net, mask, version = packed_ip_net
return (
SQL("(network_address & ?) == ?", (mask, net))
& (model.netmask >= mask)
& (model.version == version)
)
def _filter_ip_net_subnets_exclusive(
model: Type[Model], packed_ip_net: Tuple[int, int, int]
):
"""
Filters ip addresses/networks contained in ip network net.
Does not includes network itself
:param model: model to apply query
:param packed_ip_net: tuple of integers
:return: peewee expression
"""
net, mask, version = packed_ip_net
return (
SQL("(network_address & ?) == ?", (mask, net))
& (model.netmask > mask)
& (model.version == version)
)
def _filter_ip_net_supernets(
model: Type[Model], packed_ip_net: Tuple[int, int, int]
):
"""
Filters ip addresses/networks that includes provided
ip address/network, including network itself
:param model: model to apply query
:param packed_ip_net: tuple of integers
:return: peewee expression
"""
net, mask, version = packed_ip_net
return (
SQL("(? & netmask) == network_address", (net,))
& (model.netmask <= mask)
& (model.version == version)
)
def _ip_search_condition(model, ip_str):
try:
ip_net = ipaddress.ip_network(ip_str)
except ValueError:
return model.ip.contains(ip_str)
else:
net, mask, version = pack_ip_network(ip_net)
if not ip_net.hostmask:
return model.ip.contains(ip_str) | _filter_ip_net_supernets(
model, (net, mask, version)
)
else:
return (
model.ip.contains(ip_str)
| _filter_ip_net_supernets(model, (net, mask, version))
| _filter_ip_net_subnets_exclusive(model, (net, mask, version))
)
def _net_search_condition(model, ip_net):
net, mask, version = pack_ip_network(ip_net)
return _filter_ip_net_supernets(
model, (net, mask, version)
) | _filter_ip_net_subnets_exclusive(model, (net, mask, version))
def _add_ip_net_args(kwargs):
if all(k in kwargs for k in ("network_address", "netmask", "version")):
ip = unpack_ip_network(
kwargs["network_address"], kwargs["netmask"], kwargs["version"]
)
kwargs["ip"] = IP.ip_net_to_string(ip)
elif "ip" in kwargs:
# FIXME remove this when all messages will contain ipaddress.* obj
ip = IP.adopt_to_ipvX_network(kwargs["ip"])
net, mask, version = pack_ip_network(ip)
kwargs.update(
{"network_address": net, "netmask": mask, "version": version}
)
kwargs["ip"] = IP.ip_net_to_string(ip)
return kwargs
def _replace_ip_with_packed_repr(args):
if isinstance(args.get("ip"), (IPv4Network, IPv6Network)):
(
args["network_address"],
args["netmask"],
args["version"],
) = pack_ip_network(args["ip"])
del args["ip"]
return args
def is_expired(expiration) -> bool:
"""Whether expiration time passed."""
return expiration != IPList.NEVER and expiration <= time.time()
class ActionType(Enum):
"""What to do with matching IPs."""
#: See :attr:`IPList.BLACK`.
DROP = "drop"
#: See :attr:`IPList.GRAY`.
CAPTCHA = "captcha"
#: See :attr:`IPList.GRAY_SPLASHSCREEN`.
SPLASHSCREEN = "splashscreen"
#: See :attr:`.IgnoreList`.
IGNORE = "ignore"
class Purpose(str, Enum):
"""IPList's purposes understood by the agent.
An analog of i360.model.firewall.ActionType but for the new
(DEF-17989) server sync case.
"""
# note: the order here defines the order ipsets enumerated
# except for cases where the explicit "purposes" list overrides it
WHITE = "white"
white = WHITE # alias
DROP = "drop"
drop = DROP # alias
SPLASHSCREEN = "splashscreen"
splashscreen = SPLASHSCREEN # alias
CAPTCHA = "captcha"
captcha = CAPTCHA # alias
def __str__(self):
return self.value
@classmethod
def listname(cls, purpose: str):
return {
Purpose.WHITE: IPList.WHITE,
Purpose.DROP: IPList.BLACK,
Purpose.CAPTCHA: IPList.GRAY,
Purpose.SPLASHSCREEN: IPList.GRAY_SPLASHSCREEN,
}[Purpose[purpose]]
class IPList(Model):
"""The main persistent storage for various IP lists."""
#: field name in things like `SynclistResponse`.
ACTION_TYPE = "action_type"
#: Do not block the IP from this list.
WHITE = "WHITE"
#: Block the IP from this list as harshly as possible.
#: DROP packets on IPTables level and return 403 if proxied on web.
BLACK = "BLACK"
#: Display CAPTCHA for HTTP requests and DROP packets for non-web ports
#: for IP from this list.
GRAY = "GRAY"
#: Display Splash screen challenge for HTTP requests for IP from this list.
GRAY_SPLASHSCREEN = "GRAY_SPLASHSCREEN"
#: Available list names, *in priorities order*.
IP_LISTS = (WHITE, BLACK, GRAY, GRAY_SPLASHSCREEN)
IP_LIST_PRIORITIES = list(enumerate(reversed(IP_LISTS)))
#: Default scope - IP changes are not propagated to other servers in the
#: same CLN group.
SCOPE_LOCAL = "local"
#: IP changes are propagated to other servers in the same CLN group.
SCOPE_GROUP = "group"
IGNORE = "IGNORE"
# constants for version field
VERSION_IP4 = IPv4
VERSION_IP6 = IPv6
# number of records to insert in one query
BATCH_SIZE = 50
class Signals:
"""Signals to inform subscribers about changes.
Sender of event is listname.
"""
# send **kwargs: ip: IPList
added = Signal()
# send **kwargs: ip: str
deleted = Signal()
# no **kwargs
cleared = Signal()
# send **kwargs: ip: IPList
updated = Signal()
#: The IP/network as a string - a plain IP or in CIDR notation.
ip = CharField(null=False)
#: The kind of the list, one of :attr:`IP_LISTS`.
listname = CharField(
null=False,
constraints=[Check("listname in ('{}')".format("','".join(IP_LISTS)))],
)
#: expiration value meaning "it never expires"
NEVER = 0
#: Timestamp, after which this record must have no effect.
#: Nullable to be consistent with previously used create table sql.
expiration = IntegerField(
default=0, null=True # 0 - never
) # null - the same :(
#: May contain the name of a 3rd party IDS if the record was imported
#: from it.
imported_from = CharField(null=True)
#: Timestamp when the record was added to the table.
ctime = IntegerField(
null=True, default=lambda: int(time.time()) # those
) # are OK
#: A counter used to increase block TTL for consecutive blocks.
deep = IntegerField(null=True)
#: Used for :attr:`WHITE` and :attr:`BLACK` lists, can be edited by admin.
comment = CharField(null=True)
#: A reference to country code and name for the IP, based on GeoDB data.
country = CharField(null=True, column_name="country_id")
#: Relevant for non-manually whitelisted ips only.
#: Should be ignored for others.
captcha_passed = BooleanField(null=False, default=False)
#: Defines whether ip is added via UI/CLI/group i.e.,
#: whether the ip blocking is requested by user explicitly.
#: It is expected to be false for global gray/splash listings sent by
#: imunify360 server.
manual = BooleanField(null=False, default=True)
#: Indicates that whitelisted IP must have access to blocked ports too.
#: Available only for :attr:`WHITE` list.
full_access = BooleanField(null=True)
#: Indicates that IP was auto-whitelisted due to successful panel login.
auto_whitelisted = BooleanField(null=True, default=False)
#: Numeric representation of the IP.
network_address = IntegerField(null=False)
#: Numeric representation of the IP mask.
netmask = IntegerField(null=False)
#: `4` for IPv4 records and `6` for IPv6.
version = IntegerField(null=False)
#: Manual changes in IP lists are propagated to other server in the same
#: CLN group if this field is set to :attr:`SCOPE_GROUP`.
scope = CharField(
null=True,
constraints=[
Check("scope in ('%s','%s')" % (SCOPE_LOCAL, SCOPE_GROUP))
],
)
class Meta:
database = instance.db
db_table = "iplist"
primary_key = CompositeKey(
"network_address", "netmask", "version", "listname"
)
schema = "resident"
class OrderBy:
@staticmethod
def expiration():
return Case(IPList.expiration, ((0, 1),), 0), IPList.expiration
@staticmethod
def action_type2listname(action_type: Optional[str]) -> str:
"""Given `action_type` string return corresponding list name.
Return :attr:`GRAY` for an unknown/missing `action_type`.
"""
return {
ActionType.DROP.value: IPList.BLACK,
ActionType.CAPTCHA.value: IPList.GRAY,
ActionType.SPLASHSCREEN.value: IPList.GRAY_SPLASHSCREEN,
ActionType.IGNORE.value: IPList.IGNORE,
None: IPList.GRAY,
}.get(action_type, IPList.GRAY)
@staticmethod
def listname2action_type(listname: str) -> str:
"""Return action_type corresponding to iplist name."""
return {
IPList.BLACK: ActionType.DROP.value,
IPList.GRAY: ActionType.CAPTCHA.value,
IPList.GRAY_SPLASHSCREEN: ActionType.SPLASHSCREEN.value,
}[listname]
@staticmethod
def get_listname_from(properties: Optional[dict]):
"""Get iplist name corresponding to properties' action_type.
Return GRAY for an unknown/missing action_type
"""
return IPList.action_type2listname(
properties.get(IPList.ACTION_TYPE)
if properties is not None
else None
)
@staticmethod
def get_expiration_from(properties: Optional[dict]) -> int:
"""Get expiration from properties
Return IPList.NEVER if property not defined"""
return (
properties.get("expiration", IPList.NEVER)
if properties
else IPList.NEVER
)
@classmethod
def get_listname_with_priority_from(cls, properties) -> Tuple[str, int]:
"""Return tuple listname and priority."""
_listname = cls.get_listname_from(properties)
for priority, listname in cls.IP_LIST_PRIORITIES:
if listname == _listname:
return listname, priority
assert 0, "can't happen" # pragma: no cover
@classmethod
def move(
cls,
ip: IPNetwork,
src: List[str],
dest: str,
expiration: int = 0,
full_access: bool = False,
):
"""Move ip from src lists to dest list,
as `move` used only in UI and CLI we add manual=True
:param ip: ip address
:param src: src lists (WHITE/BLACK/GRAY/GRAY_SPLASHSCREEN)
:param dest: dst list (WHITE/BLACK)
:param expiration: IPs TTL. 0 means permanent
:param full_access: access to all ports
:return int: items moved
"""
assert dest not in (
cls.GRAY,
cls.GRAY_SPLASHSCREEN,
), "Move to GRAY list is not supported"
net, mask, version = pack_ip_network(ip)
(max_expiration,) = [
rec.expiration
for rec in IPList.select(fn.MAX(IPList.expiration)).where(
IPList.listname.in_(src),
IPList.network_address == net,
IPList.netmask == mask,
IPList.version == version,
)
]
IPList.delete().where(
IPList.listname.in_(src),
IPList.network_address == net,
IPList.netmask == mask,
IPList.version == version,
IPList.expiration != max_expiration,
).execute()
q = IPList.update(
listname=dest,
expiration=expiration,
full_access=full_access,
manual=True,
captcha_passed=False,
).where(IPList.listname.in_(src))
q = q.where(
IPList.network_address == net,
IPList.netmask == mask,
IPList.version == version,
)
rv = q.execute()
for listname in src:
cls.Signals.deleted.send(listname, ip=ip)
cls.Signals.added.send(dest, ip=cls.get(ip=ip))
return rv
@classmethod
def create(cls, **kwargs):
"""
:param kwargs:
:raises: IntegrityError
:return: model instance
"""
inst = super(IPList, cls).create(**_add_ip_net_args(kwargs))
cls.Signals.added.send(inst.listname, ip=inst)
return inst
@classmethod
def delete_from_list(
cls,
ip: Union[str, IPv4Network, IPv6Network],
listname: List[str],
manual: bool = None,
):
"""Delete ip from lists if exists
Return number of deleted records.
"""
if listname is None:
raise ValueError("listname should not be None")
q = IPList.delete().where(IPList.listname.in_(listname))
if manual is not None:
q = q.where(cls.manual == manual)
if isinstance(ip, str):
q = q.where(IPList.ip == ip)
else:
net, mask, ver = pack_ip_network(ip)
q = q.where(
IPList.network_address == net,
IPList.netmask == mask,
IPList.version == ver,
)
rows_deleted = q.execute()
if rows_deleted:
for lst in listname:
cls.Signals.deleted.send(lst, ip=ip)
return rows_deleted
@classmethod
def clean_lists(cls, listnames):
num_deleted = (
IPList.delete().where(IPList.listname.in_(listnames)).execute()
)
if num_deleted:
for listname in listnames:
cls.Signals.cleared.send(listname)
return num_deleted
@classmethod
def cleanup_expired_from_bglist(cls, num_days):
"""
Removes obsoleted graylist/splashscreen+blacklist[manual=False] IPs.
:param num_days: expired more than num_days ago
:return: int rows deleted
"""
expiration_ts = time.time() - timedelta(days=num_days).total_seconds()
graylist_ip_is_expired = (cls.listname == cls.GRAY) & cls.is_expired(
expiration_ts
)
graysplash_ip_is_expired = (
cls.listname == cls.GRAY_SPLASHSCREEN
) & cls.is_expired(expiration_ts)
blacklist_ip_is_expired = (
(cls.listname == cls.BLACK)
& (cls.manual == False)
& cls.is_expired(expiration_ts) # noqa: E712
)
num_deleted = (
cls.delete()
.where(
graylist_ip_is_expired
| graysplash_ip_is_expired
| blacklist_ip_is_expired
)
.execute()
)
return num_deleted
@classmethod
def delete_expired(
cls,
listname: str = None,
ip: IPNetwork = None,
):
clauses = cls.is_expired()
if listname:
clauses &= cls.listname == listname
if ip:
net, mask, version = pack_ip_network(ip)
clauses &= (
(IPList.network_address == net)
& (IPList.netmask == mask)
& (IPList.version == version)
)
deleted = cls.delete().where(clauses).execute()
return deleted
@classmethod
def _fetch_query(
cls,
listnames,
group_by=None,
having=None,
by_ip=None,
by_country_code=None,
by_comment=None,
manual=None,
):
assert isinstance(listnames, list)
# gray/splash lists do not support comment
if by_comment is not None:
listnames = [
ln
for ln in listnames
if ln not in [cls.GRAY, cls.GRAY_SPLASHSCREEN]
]
q = (
IPList.select(
IPList.ip,
IPList.listname,
IPList.expiration,
IPList.imported_from,
IPList.ctime,
IPList.deep,
IPList.comment,
IPList.country,
IPList.manual,
IPList.full_access,
IPList.auto_whitelisted,
IPList.network_address,
IPList.netmask,
IPList.version,
fn.ifnull(IPList.scope, cls.SCOPE_LOCAL).alias("scope"),
)
.join(Country, JOIN.LEFT_OUTER, on=(IPList.country == Country.id))
.where(IPList.listname.in_(listnames))
.where(~IPList.is_expired())
.order_by(IPList.ip)
)
# filter by optional args
if group_by is not None:
q = q.group_by(group_by)
if having is not None:
q = q.having(having)
if manual is not None:
q = q.where(cls.manual == manual)
if by_ip:
net = is_net(by_ip)
search_condition, net = (
(_net_search_condition, net)
if net
else (_ip_search_condition, by_ip)
)
q = q.where(search_condition(cls, net))
if by_country_code:
q = q.where(Country.code == by_country_code)
if by_comment is not None:
q = q.where(IPList.comment.contains(by_comment))
return q
@classmethod
def fetch_count(cls, listnames: List[str], **filter_args):
return cls._fetch_query(listnames, **filter_args).count()
@classmethod
def fetch_for_group_sync(cls, ips: List[IPNetwork]):
"""
cannot use 'in' operator, since peewee's @hybrid_property can not
compute lazy expression for 'IPList.ip_network' property
"""
q = IPList.select().where(
IPList.scope == IPList.SCOPE_GROUP,
)
result = []
for chunk in split_for_chunk(ips, int(CHUNK_SIZE_SQL_QUERY / 4)):
expressions = []
for ip in chunk:
net, mask, version = pack_ip_network(ip)
expressions.append(
(IPList.network_address == net)
& (IPList.netmask == mask)
& (IPList.version == version)
)
clause = reduce(ior, expressions)
result += list(q.where(clause))
return result
@classmethod
def fetch(
cls,
listnames: List[str],
offset=None,
limit=None,
order_by=None,
exclude_fields=frozenset(),
**filter_args,
):
"""
:return tuple: (max count, list of dict)
"""
q = cls._fetch_query(listnames, **filter_args)
if offset is not None:
q = q.offset(offset)
if limit is not None:
q = q.limit(limit)
if order_by is not None:
purpose_order = [
order for order in order_by if order.column_name == "purpose"
]
others_order = [
order for order in order_by if order.column_name != "purpose"
]
orders = []
# we need order records by purpose - it's mean order by
# listname priority in some custom order,
# see cls.IP_LIST_PRIORITIES.
# but our general solution `apply_order_by`
# doesn't work in this case, because it's relay on column name in
# model, so construct `order by` part of sql query in next way
for order in purpose_order:
orders.append(
cls.list_priority()
if order.desc
else cls.list_priority().desc()
)
for order in others_order:
nodes = ApplyOrderBy.get_nodes(
cls, order.column_name.split(".")
)
for node in nodes:
orders.append(node.desc() if order.desc else node)
q = q.order_by(*orders)
# use model_to_dict to convert a model instance
# to a dict recursively (with foreign keys)
rows = []
for row in q:
entry = model_to_dict(row, exclude=exclude_fields)
if entry.get("country"): # for backward compatibility
entry["country"] = model_to_dict(Country.get(id=row.country))
rows.append(entry)
return rows
@classmethod
def get_field(cls, field, default=None, **kwargs):
"""Return matching row's field value or `default` if not found"""
assert kwargs, "provides kwargs to find by them"
try:
return getattr(cls.get(**kwargs), field)
except DoesNotExist:
return default
@classmethod
def fetch_non_expired_query(cls, listname, full_access=None, version=None):
q = cls.select(cls.ip, cls.expiration).where(
(~cls.is_expired()) & (cls.listname == listname)
)
if full_access is not None:
clause = cls.full_access == full_access
if not full_access:
clause |= cls.full_access.is_null()
q = q.where(clause)
if version is not None:
q = q.where(cls.version == version)
return q
@classmethod
def fetch_non_expired(cls, listname, full_access=None, version=None):
q = cls.fetch_non_expired_query(listname, full_access, version)
try:
yield from q.dicts().iterator()
except RuntimeError:
return
@classmethod
def fetch_ipnetwork_list(
cls, listname=BLACK, manual=None
) -> Iterable[IPNetwork]:
"""
Fetch listname the most efficient (though experiementally found)
way possible.
"""
q = cls.select(IPList.network_address, IPList.netmask, IPList.version)
q = q.where((IPList.listname == listname) & ~IPList.is_expired())
if manual is not None:
q = q.where(cls.manual == manual)
for row in q.dicts():
yield unpack_ip_network(
row["network_address"], row["netmask"], row["version"]
)
@classmethod
def get_listname(cls, ip):
try:
return IPList.get(ip=ip).listname
except DoesNotExist:
pass
@classmethod
def is_expirable(cls):
"""The result has to be passed to .where(...)"""
return (cls.expiration != 0) & ~(cls.expiration >> None)
@classmethod
def is_expired(cls, expiration_ts=None):
"""The result has to be passed to .where(...)"""
if expiration_ts == IPList.NEVER:
return cls.is_expirable()
return cls.is_expirable() & (
cls.expiration <= int(expiration_ts or time.time())
)
def lives_longer(self, expiration):
"""Whether the ip record lives longer than given *expiration* time."""
return expiration != IPList.NEVER and (
self.expiration == IPList.NEVER or self.expiration > expiration
)
@classmethod
def lives_longer_prop(cls, properties, expiration):
"""Whether the properties for ip lives longer than given *expiration*
time.
"""
return expiration != IPList.NEVER and (
properties.get("expiration", IPList.NEVER) == IPList.NEVER
or properties.get("expiration") > expiration
if properties
else True
)
def lives_less(self, expiration):
"""Whether the ip record lives less than given *expiration* time."""
return self.expiration != IPList.NEVER and (
expiration == IPList.NEVER or self.expiration < expiration
)
def subnet_of(self, net, mask, version):
"""Analog of 3.7+ self.ip_network.subnet_of"""
a = self.ip_network
b = unpack_ip_network(net, mask, version)
return (
b.version == a.version
and b.network_address <= a.network_address
and b.broadcast_address >= a.broadcast_address
)
def update_properties(self, expiration, deep, listname, manual=None):
"""
Update blocking properties
:return tuple: real expiration
"""
assert expiration is not None, "'expiration' must not be None"
assert deep is not None, "'deep' must not be None"
# (x or 0) is to prevent
# TypeError: unorderable types: NoneType() < int()
self.expiration = max(self.expiration or 0, expiration)
self.deep = max(self.deep or 0, deep)
primary_key_changed = self.listname != listname
self.listname = listname
if manual is not None:
self.manual = manual
# if we change part of primary key (listname) we should create new
# record instead just save existing.
self.save(force_insert=primary_key_changed)
self.Signals.updated.send(listname, ip=self)
return {
"expiration": self.expiration,
"deep": self.deep,
}
@classmethod
def effective_list(cls, ip: IPNetwork) -> Optional[str]:
"""Return the name of highest priority list that contains the *ip*.
Return None if *ip* not in any list or record expired."""
net, mask, version = pack_ip_network(ip)
q = (
cls.select(
cls.listname,
cls.list_priority().alias("priority"),
)
.where(
~cls.is_expired(),
cls.network_address == net,
cls.netmask == mask,
cls.version == version,
)
.order_by(SQL("priority").desc())
.limit(1)
)
for r in q:
return r.listname
return None
@classmethod
def blacklist_graylisted_on_captcha_dos_alert(
cls, ip, expiration, comment
):
"""Update ip lists on CaptchaDosAlert atomically."""
with cls._meta.database.atomic():
return cls._blacklist_graylisted_on_captcha_dos_alert(
ip, expiration, comment
)
@classmethod
def _blacklist_graylisted_on_captcha_dos_alert(
cls, ip, expiration, comment
):
"""Update ip lists on CaptchaDosAlert.
Spec [1]:
if search(ip, "WHITE"): # should not really happen
return
existing_black_supernets = search(ip, "BLACK")
if any(n.expiration >= expiration for n in existing_black_supernets):
# should not really happen
return
existing = search_exactly(ip, "BLACK")
if existing.manual:
# Do nothing if already added manually
return
else:
# exact match with less expiration, remove it
# and replace with new expiration later
remove(ip, "BLACK")
# it can really exist only in GRAY list
for listname in ["GRAY", "GRAY_SPLASHSCREEN", "IGNORE"]:
existing = search_exactly(ip, listname)
if existing and existing.expiration <= expiration:
remove(ip, listname)
add(ip, "BLACK", expiration)
[1]: https://gerrit.cloudlinux.com/#/c/61260/14/src/handbook/message_processing/local_captcha_dos.py
:param ip: attackers ip
:param expiration: when record will expired
:return: Union[Dict, Exception]
""" # noqa
network, mask, version = pack_ip_network(ip)
unblocklist = [] # existing ip records (GRAY most likely)
# NOTE: include expired records to avoid conflict on insert
# and to cleanup sooner
supernets = IPList.find_closest_ip_nets(ip, expiration=1)
for net in supernets:
if net.listname == IPList.WHITE and net.lives_longer(time.time()):
# found non-expired whitelist record unexpectedly
return Reject(
"Don't blacklist whitelisted ips [%s] on CaptchaDosAlert"
% (ip,)
)
if (
net.listname == IPList.BLACK
and net.lives_longer(time.time())
and net.manual
):
# a non-expired manual BLACK record tramps CaptchaDos
# (even if the latter may expire later)
return Reject(
"Don't blacklist manually blacklisted ips [%s] on CaptchaDosAlert" # noqa
% (ip,)
)
if net.listname == IPList.BLACK and (
net.lives_longer(expiration) or not net.lives_less(expiration)
): # >=
# NOTE: assume expiration > now
# no need to blacklist
return Reject(
"%s is already blacklisted for long enough time: %s >= %s"
% (
ip,
net.expiration,
expiration,
)
)
# exact match by ip and the record doesn't outlive alert
if net.ip_network == ip and not net.lives_longer(expiration):
# <=
unblocklist.append((net.ip, net.listname))
# remove ip records unconditionally
for ip_, listname in unblocklist:
rows_deleted = cls.delete_from_list(ip_, [listname])
assert rows_deleted
# remove subnets from IGNORE list
is_subnet = (
IgnoreList.network_address.bin_and(mask) == network,
IgnoreList.netmask >= mask,
IgnoreList.version == version,
)
ignore_subnets = (
IgnoreList.select(IgnoreList.ip).where(*is_subnet).execute()
)
unblocklist += [(r.ip, IPList.IGNORE) for r in ignore_subnets]
IgnoreList.delete().where(*is_subnet).execute()
if unblocklist:
logger.info(
"Removed %s from %s lists", ip, [L for _, L in unblocklist]
)
# add ip record
cls.create(
ip=ip,
listname=IPList.BLACK,
expiration=expiration,
comment=comment,
manual=False,
)
logger.info("Put %s on the BLACK list", ip)
return dict(
blocklist={
# primary key -> other ip properties
(ip, IPList.BLACK): dict(expiration=expiration)
},
unblocklist=[(ip, listname) for _, listname in unblocklist],
)
@property
def ip_network(self):
return unpack_ip_network(
self.network_address, self.netmask, self.version
)
@classmethod
def find_closest_ip_nets(
cls,
ip: Union[ipaddress.IPv4Network, ipaddress.IPv6Network],
listname: List[str] = None,
limit: int = None,
*,
expiration=None,
manual=None,
):
"""
Returns all supernets containing given network (*ip*)
that are not expired by *expiration* time.
:param ip: ip network to lookup
:param listname: list of listnames
:param limit: number of supernets to return
:param expiration: seconds since the epoch or None
None means "use the current time"
:return: list of matching supernets, ordered by netmask
(desc: from smallest to largest)
"""
net, mask, version = pack_ip_network(ip)
q = cls.select().where(
_filter_ip_net_supernets(cls, (net, mask, version))
)
q = q.where(~cls.is_expired(expiration))
if listname is not None:
q = q.where(cls.listname.in_(listname))
q = q.order_by(cls.netmask.desc())
if manual is not None:
q = q.where(cls.manual == manual)
if limit is not None:
q = q.limit(limit)
return list(q)
@classmethod
def list_priority(cls):
return Case(
None,
[
(cls.listname == listname, priority)
for priority, listname in cls.IP_LIST_PRIORITIES
],
)
@classmethod
def filter_ips_has_supernets(
cls,
networks: Dict[IPNetwork, Dict[str, Any]],
listnames: Optional[List] = None,
expiration_ts=None,
):
"""Yield networks which has lasting supernets with higher priority in
db.
If listname isn't provided ignore priorities (for unblock).
Implemented to solve performance issue: DEF-15123
"""
_TempIPList.fill(networks)
q = IPList.select(
_TempIPList.network_address,
_TempIPList.netmask,
_TempIPList.version,
).where(
(
((IPList.version == 4) & (IPList.netmask != IPV4_HOST_MASK))
| ((IPList.version == 6) & (IPList.netmask != IPV6_HOST_MASK))
)
& (~IPList.is_expired(expiration_ts))
)
if listnames is None:
q = q.where(IPList.listname != IPList.WHITE)
else:
q = q.where(IPList.listname.in_(listnames))
q = q.join(
_TempIPList,
on=(
(IPList.version == _TempIPList.version)
& (
_TempIPList.network_address.bin_and(IPList.netmask)
== IPList.network_address
)
& (IPList.netmask < _TempIPList.netmask)
& (
(
(
(IPList.expiration != IPList.NEVER)
& (IPList.expiration >= _TempIPList.expiration)
)
| (IPList.expiration == IPList.NEVER)
)
| (_TempIPList.expiration == IPList.NEVER)
)
& (
True
if listnames is None
else (cls.list_priority() >= _TempIPList.priority)
)
),
)
# note: using distinct may affect performance
# see DEF-15964
for network in set(q.tuples()):
yield unpack_ip_network(*network)
@classmethod
def find_ips_with_later_expiration(
cls, ips: Dict[IPNetwork, Dict[str, Any]]
):
"""Yield exact match of ips which already recorded in db with the same
priority and later expiration (so we don't need to add them)
Implemented to solve performance issue: DEF-15123
"""
_TempIPList.fill(ips)
q = (
cls.select(
_TempIPList.network_address,
_TempIPList.netmask,
_TempIPList.version,
)
.where((~IPList.is_expired()) & (~IPList.manual))
.join(
_TempIPList,
on=(
(IPList.version == _TempIPList.version)
& (IPList.network_address == _TempIPList.network_address)
& (IPList.netmask == _TempIPList.netmask)
& (IPList.list_priority() == _TempIPList.priority)
& (
(
(IPList.expiration != IPList.NEVER)
& (IPList.expiration >= _TempIPList.expiration)
)
| (IPList.expiration == IPList.NEVER)
)
& (_TempIPList.expiration != IPList.NEVER)
),
)
)
# note: using distinct may affect performance
# see DEF-15964
for network in set(q.tuples()):
yield unpack_ip_network(*network)
@classmethod
def find_ip_subnets_with_less_priority(
cls, ips: Dict[IPNetwork, Dict[str, Any]]
):
"""Yield
- subnet (include self) with less priority
(listname and less expiration),
- listname of the subnet
- should_unblock which is True if subnet is exact blocked
network with same listname
Implemented to solve performance issue: DEF-15123
"""
_TempIPList.fill(ips)
q = cls.select(
IPList.network_address,
IPList.netmask,
IPList.version,
IPList.listname,
(
(_TempIPList.listname != IPList.listname)
| ~(
(_TempIPList.netmask == IPList.netmask)
& (_TempIPList.network_address == IPList.network_address)
& (IPList.version == _TempIPList.version)
)
),
).join(
_TempIPList,
on=(
(IPList.version == _TempIPList.version)
& (
IPList.network_address.bin_and(_TempIPList.netmask)
== _TempIPList.network_address
)
& (IPList.netmask >= _TempIPList.netmask)
& (IPList.list_priority() <= _TempIPList.priority)
& (
(
(IPList.expiration != IPList.NEVER)
& (IPList.expiration <= _TempIPList.expiration)
)
| (_TempIPList.expiration == IPList.NEVER)
)
),
)
for (
network_address,
netmask,
version,
listname,
should_unblock,
) in q.tuples():
yield unpack_ip_network(
network_address, netmask, version
), listname, should_unblock
@classmethod
def find_lists(
cls,
ip: IPNetwork,
listname: List[str],
*,
manual=None,
expiration=None,
):
"""
Returns all lists containing given network (*ip*)
:param ip: ip network to lookup
:param listname: list of listnames
:return: names of list
"""
net, mask, version = pack_ip_network(ip)
q = cls.select(cls.listname).where(
(cls.network_address == net)
& (cls.netmask == mask)
& (cls.version == version)
)
q = q.where(~cls.is_expired(expiration))
if listname is not None:
q = q.where(cls.listname.in_(listname))
q = q.order_by(cls.listname.desc())
if manual is not None:
q = q.where(cls.manual == manual)
return list(q)
@classmethod
def get(cls, *query, **kwargs):
return super().get(*query, **_replace_ip_with_packed_repr(kwargs))
@classmethod
def get_or_create(cls, **kwargs):
return super().get_or_create(**_replace_ip_with_packed_repr(kwargs))
@classmethod
def create_or_get(cls, **kwargs):
return super().create_or_get(**_replace_ip_with_packed_repr(kwargs))
@classmethod
def find_net_members(
cls,
ip: IPNetwork,
listname: List[str] = None,
*,
include_itself=False,
manual: bool = None,
expired_by=None,
) -> List[Tuple[IPNetwork, str, int]]:
"""Return ip_network objects containing all
*ip* entries expired by *expired_by* from lists *listname*
which are members of net *ip* including itself if *include_itself*
:param ip: network to lookup members for
:param listname: list name
:param include_itself: whether to include ip itself as a subnet
[default: False]
:param expired_by: expiry date as "seconds since epoch"
return entries expired by given *expired_by* timestamp
- IPList.NEVER :: return all entries (regardless expiration)
- None :: return non-expired entries
"""
net, mask, version = pack_ip_network(ip)
q = cls.select(
cls.network_address, cls.netmask, cls.listname, cls.expiration
).where(
_filter_ip_net_subnets(cls, (net, mask, version))
if include_itself
else _filter_ip_net_subnets_exclusive(cls, (net, mask, version))
)
if expired_by is None:
q = q.where(~cls.is_expired())
elif expired_by != IPList.NEVER:
q = q.where(cls.is_expired(expired_by))
if listname is not None:
q = q.where(cls.listname.in_(listname))
if manual is not None:
q = q.where(cls.manual == manual)
return [
(unpack_ip_network(net, mask, version), listname, e)
for net, mask, listname, e in q.tuples()
]
@classmethod
def lists_with_less_or_equal_priorities(
cls, listname: str
) -> Sequence[str]:
"""Return list of iplists with less priority than list from given
properties"""
index = cls.IP_LISTS.index(listname)
return cls.IP_LISTS[index:]
@classmethod
def lists_with_greater_or_equal_priorities(
cls, listname: str
) -> Sequence[str]:
"""Return list of iplists with greater priority than list from given
listname"""
index = cls.IP_LISTS.index(listname)
return cls.IP_LISTS[: index + 1]
@classmethod
def remove(cls, ip: IPNetwork, listname: str, manual=None):
network, mask, version = pack_ip_network(ip)
query = cls.delete().where(
cls.listname == listname,
cls.network_address == network,
cls.netmask == mask,
cls.version == version,
)
if manual is not None:
query = query.where(cls.manual == manual)
return query.execute()
@classmethod
def block_many(cls, to_block: List[Dict]):
for idx in range(0, len(to_block), cls.BATCH_SIZE):
cls.insert_many(to_block[idx : idx + cls.BATCH_SIZE]).on_conflict(
# consider existing record should be replaced silently
"REPLACE"
).execute()
@classmethod
def remove_many(
cls, ips: Iterable[Tuple[IPNetwork, Dict]]
) -> List[Tuple[IPNetwork, str]]:
"""
Remove *ips* that are not manual from the [iplist] table.
Return ips to unblock.
"""
_TempIPList.fill(ips)
primary_key_sql = SQL(
"({})".format(", ".join(cls._meta.primary_key.field_names))
)
primary_key = (
getattr(cls, field) for field in cls._meta.primary_key.field_names
)
to_remove = (
cls.select(*primary_key)
.where(~IPList.manual)
.join(
_TempIPList,
on=(
(IPList.version == _TempIPList.version)
& (IPList.network_address == _TempIPList.network_address)
& (IPList.netmask == _TempIPList.netmask)
& (IPList.listname == _TempIPList.listname)
),
)
)
to_unblock = []
for network_address, netmask, version, listname in to_remove.tuples():
to_unblock.append(
(
unpack_ip_network(network_address, netmask, version),
listname,
)
)
cls.delete().where(primary_key_sql.in_(to_remove)).execute()
return to_unblock
class _TempIPList(Model):
network_address = IntegerField(null=False)
netmask = IntegerField(null=False)
version = IntegerField(null=False)
listname = CharField(null=False)
priority = IntegerField(null=False)
expiration = IntegerField(
default=0, null=True # 0 - never
) # null - the same :(
class Meta:
database = instance.db
db_table = "tmp_iplist"
primary_key = CompositeKey(
"network_address", "netmask", "version", "listname"
)
@classmethod
def _create(cls):
cls._meta.database.execute_sql(
"""
CREATE TEMPORARY TABLE IF NOT EXISTS tmp_iplist
(
network_address INT, netmask INT, version INT,
listname TXT VARCHAR(255) NOT NULL CHECK
(listname in ('WHITE','BLACK','GRAY','GRAY_SPLASHSCREEN')),
priority INT,
expiration INT,
PRIMARY KEY (network_address, netmask, version, listname))
"""
)
@classmethod
def _clear(cls):
cls.delete().execute()
@classmethod
def fill(cls, ips: Union[Dict, Iterable]):
ips = ips.items() if hasattr(ips, "items") else ips # type: ignore
with timeit("prepare tmp_iplist to fill", logger):
cls._create()
cls._clear()
data = [
dict(
zip(
[
"network_address",
"netmask",
"version",
"listname",
"priority",
"expiration",
],
[
*pack_ip_network(ip),
*IPList.get_listname_with_priority_from(prop),
IPList.get_expiration_from(prop),
],
)
)
for ip, prop in ips
]
if data:
with timeit("fill tmp_iplist", logger):
batch_size = 150
for idx in range(0, len(data), batch_size):
cls.insert_many(data[idx : idx + batch_size]).execute()
class LastSynclist(Model):
"""Used to track how up-to-date are lists synced from Correlation server."""
IP, HASH = "ip", "hash"
#: When the last synclist happened.
timestamp = FloatField(null=True, default=0)
#: Synclist type. Currently only `ip` is supported.
name = CharField(null=False, primary_key=True)
class Meta:
database = instance.db
db_table = "last_synclist"
schema = "resident"
@classmethod
def update_timestamp(cls, name):
cls.update(timestamp=time.time()).where(cls.name == name).execute()
@classmethod
def get_timestamp(cls, name):
obj, _ = cls.get_or_create(name=name, defaults={"timestamp": 0})
return obj.timestamp
class WhitelistedCrawler(Model):
"""
Crawlers for which local alerts must not add IP
to the :attr:`IPList.GRAY` list.
"""
class Meta:
database = instance.db
db_table = "whitelisted_crawlers"
id = PrimaryKeyField()
#: The name of the crawler.
description = TextField(null=False)
@classmethod
def add(cls, description, domains):
with instance.db.atomic():
inserted_id = cls.insert(description=description).execute()
for d in domains:
WhitelistedCrawlerDomain.create(
crawler_id=inserted_id, domain=d
)
@classmethod
def fetch(cls, limit, offset):
crawlers_query = (
cls.select().order_by(cls.description).limit(limit).offset(offset)
)
domains_query = WhitelistedCrawlerDomain.select()
crawlers_with_domains_query = prefetch(crawlers_query, domains_query)
result = []
max_count = crawlers_query.count(clear_limit=True)
for crawler in crawlers_with_domains_query:
item = {
"id": crawler.id,
"description": crawler.description,
"domains": [d.domain for d in crawler.domains],
}
result.append(item)
return max_count, result
class WhitelistedCrawlerDomain(Model):
"""Domain names used to check if IP is a :class:`WhitelistedCrawler`."""
class Meta:
database = instance.db
db_table = "whitelisted_crawler_domains"
id = PrimaryKeyField()
crawler = ForeignKeyField(
WhitelistedCrawler,
null=False,
on_delete="CASCADE",
related_name="domains",
)
#: Domain name used to check if IP is a known whitelisted crawler,
#: via reverse & forward DNS lookup.
domain = TextField(null=False)
class RemoteProxyGroup(Model):
"""Groups multiple remote proxies together with common data."""
#: Proxy Groups added by admin.
MANUAL = "manual"
#: Proxy Groups added internally.
IMUNIFY360 = "imunify360"
#: The name of the Proxy Group.
name = CharField(null=False)
#: Who added the group,
#: must be either :attr:`MANUAL` or :attr:`IMUNIFY360`.
source = CharField(
null=False,
constraints=[
Check("source in ('{}', '{}')".format(MANUAL, IMUNIFY360))
],
)
#: Disabled groups are not effective (not added to IPSets, WebShield).
enabled = BooleanField(null=False, default=True)
class Meta:
database = instance.db
db_table = "remote_proxy_group"
indexes = ((("name", "source"), True),)
@classmethod
def set_enabled(cls, name: str, source: str, enabled: bool) -> bool:
"""Set group's enabled status.
Group is identified by name and source. Returns True if enabled
status has changed, False otherwise (it was a noop)."""
group = RemoteProxyGroup.get(name=name, source=source)
if group.enabled == enabled:
return False
group.enabled = enabled
group.save()
return True
class RemoteProxy(Model):
"""Remote Proxy networks in a group."""
group = ForeignKeyField(RemoteProxyGroup, null=False)
#: A network as a string in CIDR notation.
network = TextField(null=False)
class Meta:
database = instance.db
db_table = "remote_proxy"
@classmethod
def list(
cls,
by_group: Optional[str],
by_source: Optional[str],
enabled: Optional[bool],
) -> List[dict]:
"""Returns a list of remote proxy networks as dicts.
Results are optionally filtered by group name, source, and enabled
status."""
q = cls.select(
RemoteProxyGroup.source,
RemoteProxyGroup.name,
RemoteProxyGroup.enabled,
cls.network,
).join(RemoteProxyGroup)
if by_group is not None:
q = q.where(RemoteProxyGroup.name == by_group)
if by_source is not None:
q = q.where(RemoteProxyGroup.source == by_source)
if enabled is not None:
q = q.where(RemoteProxyGroup.enabled == enabled)
return list(q.order_by(RemoteProxyGroup.name).dicts())
@classmethod
def add_many(cls, name: str, source: str, networks: List[str]):
"""Adds networks to a list of remote proxy in group name, source."""
group, _ = RemoteProxyGroup.get_or_create(name=name, source=source)
for net in networks:
# convert to common representation
net = IP.ip_net_to_string(ipaddress.ip_network(net))
proxy = RemoteProxy(network=net, group_id=group.id)
proxy.save()
@classmethod
def delete_networks(cls, source: str, networks: List[str]):
"""Deletes networks from remote proxy lists.
Only networks coming from groups with given source are deleted."""
# convert to common representation
networks = [
IP.ip_net_to_string(ipaddress.ip_network(net)) for net in networks
]
deleted = []
# delete requested networks
q = (
RemoteProxy.select()
.join(RemoteProxyGroup)
.where(RemoteProxy.network << networks)
.where(RemoteProxyGroup.source == source)
)
for proxy in list(q):
deleted.append(proxy.network)
proxy.delete_instance()
# delete groups having no members
q = (
RemoteProxyGroup.select(RemoteProxyGroup)
.join(RemoteProxy, JOIN.LEFT_OUTER)
.group_by(RemoteProxyGroup)
.having(fn.COUNT(RemoteProxy.id) == 0)
)
for group in list(q):
group.delete_instance()
return deleted
class IgnoreList(Model):
"""
IP addresses from this list are not blocked by firewall. However,
they still can be placed to other lists by either server or local
events or by user request.
"""
#: The IP/network as a string - a plain IP or in CIDR notation.
ip = CharField(null=False)
#: Numeric representation of the IP.
network_address = IntegerField(null=False)
#: Numeric representation of the IP mask.
netmask = IntegerField(null=False)
#: `4` for IPv4 records and `6` for IPv6.
version = IntegerField(null=False)
class Meta:
database = instance.db
db_table = "ignore_list"
primary_key = CompositeKey("network_address", "netmask", "version")
schema = "resident"
@classmethod
def create(cls, **kwargs):
"""
:param kwargs:
:raises: IntegrityError
:return: model instance
"""
return super().create(**_add_ip_net_args(kwargs))
@classmethod
def create_or_get(cls, **kwargs):
return super().create_or_get(**_replace_ip_with_packed_repr(kwargs))
@classmethod
def get(cls, *query, **kwargs):
return super().get(*query, **_replace_ip_with_packed_repr(kwargs))
@classmethod
def get_or_create(cls, **kwargs):
return super().get_or_create(**_replace_ip_with_packed_repr(kwargs))
@classmethod
def subnets(cls, supernet: IPNetwork) -> Iterable[IPNetwork]:
address, mask, version = pack_ip_network(supernet)
q = cls.select().where(
(cls.network_address.bin_and(mask)) == address,
cls.netmask >= mask,
cls.version == version,
)
for row in q:
yield unpack_ip_network(
row.network_address, row.netmask, row.version
)
@classmethod
def remove(cls, to_delete: List[IPNetwork]):
unique = set(to_delete)
for item in unique:
address, mask, version = pack_ip_network(item)
cls.delete().where(
cls.network_address == address,
cls.netmask == mask,
cls.version == version,
).execute()
class BlockedPort(Model):
"""Port blocking configuration.
Effective when `FIREWALL.port_blocking_mode == ALLOW`.
"""
#: The port to be blocked.
port = IntegerField(null=False)
#: The protocol to be blocked - either `tcp` or `udp` or `all`.
proto = CharField(
null=False,
constraints=[Check(f"proto in ('{TCP}', '{UDP}', '{ALL}')")],
)
#: Comment set by admin.
comment = CharField(null=True)
class Meta:
database = instance.db
db_table = "blocked_port"
schema = "resident"
indexes = (
# create an unique on port/proto
(("port", "proto"), True),
)
@classmethod
def _add_filter_args(
cls, q, by_comment=None, by_ip=None, by_country_code=None
):
if by_comment is not None:
q = q.where(
cls.comment.contains(by_comment)
| IgnoredByPort.comment.contains(by_comment)
)
if by_ip is not None:
q = q.where(_ip_search_condition(IgnoredByPort, by_ip))
if by_country_code is not None:
q = q.where(Country.code == by_country_code)
return q
@classmethod
def fetch_count(cls, **filter_args):
q = (
cls.select(cls.id)
.distinct()
.join(IgnoredByPort, JOIN.LEFT_OUTER)
.join(
Country,
JOIN.LEFT_OUTER,
on=(IgnoredByPort.country == Country.id),
)
)
q = cls._add_filter_args(q, **filter_args)
return q.count()
@classmethod
def fetch(cls, limit=50, offset=0, **filter_args):
q = (
cls.select(
cls.id,
cls.port,
cls.proto,
cls.comment,
IgnoredByPort.ip,
IgnoredByPort.comment.alias("ip_comment"),
)
.join(IgnoredByPort, JOIN.LEFT_OUTER)
.join(
Country,
JOIN.LEFT_OUTER,
on=(IgnoredByPort.country == Country.id),
)
)
q = cls._add_filter_args(q, **filter_args)
q = q.order_by(cls.port, cls.proto, IgnoredByPort.ip)
def group_key(row):
return {
key: row[key] for key in ("id", "port", "proto", "comment")
}
result = []
for port_proto, ips in itertools.groupby(q.dicts(), key=group_key):
ignored_ips = [
{"ip": ip["ip"], "comment": ip["ip_comment"]}
for ip in ips
if ip["ip"] is not None
]
port_proto["ips"] = ignored_ips
result.append(port_proto)
return result[offset : offset + limit]
class IgnoredByPort(Model):
"""Ignored IPs for ports blocked via :class:`BlockedPort`."""
port_proto = ForeignKeyField(
BlockedPort, null=False, on_delete="CASCADE", related_name="ips"
)
#: The IP/network as a string - a plain IP or in CIDR notation.
ip = CharField(null=False)
#: Comment set by admin.
comment = CharField(null=True)
#: Numeric representation of the IP.
network_address = IntegerField(null=False)
#: Numeric representation of the IP mask.
netmask = IntegerField(null=False)
#: `4` for IPv4 records and `6` for IPv6.
version = IntegerField(null=False)
#: A reference to country code and name for the IP, based on GeoDB data.
country = CharField(null=True, column_name="country_id")
class Meta:
database = instance.db
db_table = "ignored_by_port_proto"
schema = "resident"
indexes = (
# create an unique on port/ip
(("port_proto", "ip"), True),
)
@classmethod
def create(cls, **kwargs):
super().create(**_add_ip_net_args(kwargs))
@classmethod
def fetch(cls, version: NumericIPVersion = None):
q = cls.select(cls, BlockedPort).join(BlockedPort)
if version is not None:
q = q.where(cls.version == version)
return q
# We hit a bug in peewee, which generates wrong sql on table creation,
# which is why create_table overridden
# Namely:
# REFERENCES "resident"."blocked_port"
# ^^^^^^^ this part should be ommited
@classmethod
def create_table(cls, safe=True, **options):
# dynamic scheme change (used in unit tests only)
schema = f"{cls._meta.schema}." * bool(cls._meta.schema)
cls._meta.database.execute_sql(
f"""
CREATE TABLE IF NOT EXISTS {schema}ignored_by_port_proto
(
"id" INTEGER NOT NULL PRIMARY KEY,
"port_proto_id" INTEGER NOT NULL,
"ip" VARCHAR(255) NOT NULL,
"comment" VARCHAR(255),
"network_address" INTEGER NOT NULL,
"netmask" INTEGER NOT NULL,
"version" INTEGER NOT NULL,
"country_id" VARCHAR(255),
FOREIGN KEY ("port_proto_id")
REFERENCES "blocked_port" ("id") ON DELETE CASCADE
)"""
)
cls._schema.create_indexes(safe=safe)
class IPListRecord(Model):
"""DB table that stores ips for given iplist_id`s."""
network_address = IntegerField(null=False)
netmask = IntegerField(null=False)
version = IntegerField(null=False)
iplist_id = IntegerField(null=False)
class Meta:
database = instance.db
db_table = "iplistrecord"
primary_key = CompositeKey(
"network_address", "netmask", "version", "iplist_id"
)
schema = "ipsetlists"
@classmethod
def fetch_ips(
db, ip_version: IPVersion, iplist_id: IPListID, chunk_size=50000
) -> Iterator[IPNetwork]:
"""
Yield ips corresponding to *ip_version*, *iplist_id*.
NOTE: It assumes exclusive access to the db until the iterator
is exhausted.
"""
# To avoid too much memory consumption, don't get all records at once
# use chunks instead.
# We can't use some guidelines from https://docs.peewee-orm.com/en/latest/peewee/querying.html#iterating-over-large-result-sets # noqa
# Using .iterator() leads to
# RuntimeError: generator raised StopIteration
# or:
# sqlite3.OperationalError: database table is locked
# Using underlying database cursor db._meta.database.execute(query)
# leads to
# peewee.OperationalError: near "(": syntax error
query = (
db.select(db.network_address, db.netmask, db.version)
.where(
(db.iplist_id == iplist_id)
& (db.version == NumericIPVersion[ip_version])
)
.tuples()
)
offset = 0
# It assumes exclusive access to the db until the iterator
# is exhausted, otherwise, an `await` somewhere far away from this code
# may break it or lead to unexpected results such as duplicating ips
while chunk := query.limit(chunk_size).offset(offset):
offset += chunk.count()
yield from starmap(unpack_ip_network, chunk)
@classmethod
def fetch_ips_count(db, ip_version: IPVersion, iplist_id: IPListID) -> int:
"""How many *ip_version* ips for *iplist_id*."""
return (
db.select()
.where(
(db.iplist_id == iplist_id)
& (db.version == NumericIPVersion[ip_version])
)
.count()
)
@classmethod
def _fetch_query(cls, purpose, by_ip):
q = cls.select(
cls,
IPListPurpose.purpose,
IPListPurpose.iplist_id.alias("id"),
).join(IPListPurpose, on=(cls.iplist_id == IPListPurpose.iplist_id))
if purpose is not None:
q = q.where(IPListPurpose.purpose == purpose)
if by_ip:
q = q.where(_net_search_condition(cls, by_ip))
return q
@classmethod
def fetch_count(cls, purpose, by_ip=None):
q = cls._fetch_query(purpose, by_ip)
return q.count()
@classmethod
def fetch(cls, purpose, by_ip, offset=None, limit=None):
q = cls._fetch_query(purpose, by_ip)
if offset is not None:
q = q.offset(offset)
if limit is not None:
q = q.limit(limit)
return [
{
"network_address": str(ip.network_address),
"netmask": str(ip.netmask),
"version": row.version,
"iplist_id": row.iplistpurpose.id,
"purpose": row.iplistpurpose.purpose,
}
for row, ip in map(
lambda row: (
row,
unpack_ip_network(
row.network_address, row.netmask, row.version
),
),
q,
)
]
@classmethod
def create(cls, **kwargs):
inst = super().create(**_add_ip_net_args(kwargs))
return inst
class IPListPurpose(Model):
"""DB table that stores "purposes" for given iplist_id`s."""
purpose = CharField(null=False)
iplist_id = IntegerField(null=False)
class Meta:
database = instance.db
db_table = "iplistpurpose"
primary_key = CompositeKey("purpose", "iplist_id")
schema = "ipsetlists"
@classmethod
def fetch_iplist_ids(
db, ip_version: IPVersion, purposes: Iterable[Purpose]
) -> Iterable[IPListID]:
"""Yield all distinct iplist_id for *ip_version* and *purposes*."""
return map(
itemgetter(0),
db.select(db.iplist_id)
.distinct()
.join(IPListRecord, on=(db.iplist_id == IPListRecord.iplist_id))
.where(
(IPListRecord.version == NumericIPVersion[ip_version])
& db.purpose.in_(list(map(str, purposes)))
)
.tuples(),
)
@classmethod
def fetch_count(db, ip_version: IPVersion, purpose: Purpose) -> int:
"""How many iplists are there for *ip_version* and *purpose*"""
return (
db.select(db.iplist_id)
.distinct()
.join(IPListRecord, on=(db.iplist_id == IPListRecord.iplist_id))
.where(
(IPListRecord.version == NumericIPVersion[ip_version])
& (db.purpose == purpose)
)
.count()
)
@staticmethod
def listname2purpose(listname: str) -> Purpose:
"""Return purpose corresponding to iplist name."""
return {
IPList.WHITE: Purpose.WHITE,
IPList.BLACK: Purpose.DROP,
IPList.GRAY: Purpose.CAPTCHA,
IPList.GRAY_SPLASHSCREEN: Purpose.SPLASHSCREEN,
}[listname]
Mini Shell