<commit_before>"""iptables helper functions.

Unlike the `firewall` module, these functions know nothing about PaaSTA and
could effectively be a third-party library. They just make working with
iptables a little bit easier.
"""
from __future__ import absolute_import
from __future__ import unicode_literals

import collections
import re

import iptc


CHAIN_REGEX = re.compile('-N ([^ ]+)$')


class Rule(collections.namedtuple('Rule', (
    'protocol',
    'src',
    'dst',
    'target',
    'matches',
))):
    """Rule representation.

    Working with iptc's rule classes directly doesn't work well, since rules
    represent actual existing iptables rules, and changes are applied
    immediately. They're also difficult to compare.
    """

    @classmethod
    def from_iptc(cls, rule):
        fields = {
            'protocol': rule.protocol,
            'src': rule.src,
            'dst': rule.dst,
            'target': rule.target.name,
            'matches': (),
        }

        for match in rule.matches:
            fields['matches'] += ((
                match.name,
                tuple((param, value) for param, value in match.parameters.items())
            ),)

        return cls(**fields)

    def to_iptc(self):
        rule = iptc.Rule()
        rule.protocol = self.protocol
        rule.src = self.src
        rule.dst = self.dst
        rule.create_target(self.target)
        for name, params in self.matches:
            match = rule.create_match(name)
            for key, value in params:
                setattr(match, key, value)
        return rule


class ChainDoesNotExist(Exception):
    pass


def all_chains():
    return {chain.name for chain in iptc.Table(iptc.Table.FILTER).chains}


def ensure_chain(chain, rules):
    """Idempotently ensure a chain exists and has an exact set of rules.

    This function creates or updates an existing chain to match the rules
    passed in.

    This function will not reorder existing rules, but any new rules are always
    inserted at the front of the chain.
    """
    try:
        current_rules = list_chain(chain)
    except ChainDoesNotExist:
        create_chain(chain)
        current_rules = set()

    for rule in rules:
        if rule not in current_rules:
            insert_rule(chain, rule)

    # The insertion stage needs "rules" to be ordered since insertion order
    # matters, but the deletion phase doesn't, so we turn it into a set to
    # speed up this loop.
    rules = set(rules)
    for rule in current_rules:
        if rule not in rules:
            delete_rule(chain, rule)


def ensure_rule(chain, rule):
    rules = list_chain(chain)
    if rule not in rules:
        insert_rule(chain, rule)


def insert_rule(chain_name, rule):
    print('adding rule to {}: {}'.format(chain_name, rule))
    chain = iptc.Chain(iptc.Table(iptc.Table.FILTER), chain_name)
    chain.insert_rule(rule.to_iptc())


def delete_rule(chain_name, rule):
    print('deleting rule from {}: {}'.format(chain_name, rule))
    chain = iptc.Chain(iptc.Table(iptc.Table.FILTER), chain_name)
    for potential_rule in chain.rules:
        if Rule.from_iptc(potential_rule) == rule:
            chain.delete_rule(potential_rule)


def create_chain(chain_name):
    print('creating chain: {}'.format(chain_name))
    iptc.Table(iptc.Table.FILTER).create_chain(chain_name)


def delete_chain(chain_name):
    print('deleting chain: {}'.format(chain_name))
    chain = iptc.Chain(iptc.Table(iptc.Table.FILTER), chain_name)
    chain.flush()
    chain.delete()


def list_chain(chain_name):
    """List rules in a chain.

    Returns a list of iptables rules, or raises ChainDoesNotExist.
    """
    table = iptc.Table(iptc.Table.FILTER)<commit_msg>Turn off autocommit when deleting rules<commit_after>"""iptables helper functions.

Unlike the `firewall` module, these functions know nothing about PaaSTA and
could effectively be a third-party library. They just make working with
iptables a little bit easier.
"""
from __future__ import absolute_import
from __future__ import unicode_literals

import contextlib
import collections
import re

import iptc


CHAIN_REGEX = re.compile('-N ([^ ]+)$')


class Rule(collections.namedtuple('Rule', (
    'protocol',
    'src',
    'dst',
    'target',
    'matches',
))):
    """Rule representation.

    Working with iptc's rule classes directly doesn't work well, since rules
    represent actual existing iptables rules, and changes are applied
    immediately. They're also difficult to compare.
    """

    @classmethod
    def from_iptc(cls, rule):
        fields = {
            'protocol': rule.protocol,
            'src': rule.src,
            'dst': rule.dst,
            'target': rule.target.name,
            'matches': (),
        }

        for match in rule.matches:
            fields['matches'] += ((
                match.name,
                tuple((param, value) for param, value in match.parameters.items())
            ),)

        return cls(**fields)

    def to_iptc(self):
        rule = iptc.Rule()
        rule.protocol = self.protocol
        rule.src = self.src
        rule.dst = self.dst
        rule.create_target(self.target)
        for name, params in self.matches:
            match = rule.create_match(name)
            for key, value in params:
                setattr(match, key, value)
        return rule


@contextlib.contextmanager
def iptables_txn(table):
    """Temporarily disable autocommit and commit at the end.

    By default, changes to iptables rules are applied immediately. In some
    cases, we want to avoid that.

    https://github.com/ldx/python-iptables#autocommit
    """
    try:
        table.autocommit = False
        yield
    except:
        raise
    else:
        table.commit()
        table.refresh()
    finally:
        table.autocommit = True


class ChainDoesNotExist(Exception):
    pass


def all_chains():
    return {chain.name for chain in iptc.Table(iptc.Table.FILTER).chains}


def ensure_chain(chain, rules):
    """Idempotently ensure a chain exists and has an exact set of rules.

    This function creates or updates an existing chain to match the rules
    passed in.

    This function will not reorder existing rules, but any new rules are always
    inserted at the front of the chain.
    """
    try:
        current_rules = list_chain(chain)
    except ChainDoesNotExist:
        create_chain(chain)
        current_rules = set()

    for rule in rules:
        if rule not in current_rules:
            insert_rule(chain, rule)

    # The insertion stage needs "rules" to be ordered since insertion order
    # matters, but the deletion phase doesn't, so we turn it into a set to
    # speed up this loop.
    rules = set(rules)
    for rule in current_rules:
        if rule not in rules:
            delete_rule(chain, rule)


def ensure_rule(chain, rule):
    rules = list_chain(chain)
    if rule not in rules:
        insert_rule(chain, rule)


def insert_rule(chain_name, rule):
    print('adding rule to {}: {}'.format(chain_name, rule))
    chain = iptc.Chain(iptc.Table(iptc.Table.FILTER), chain_name)
    chain.insert_rule(rule.to_iptc())


def delete_rule(chain_name, rule):
    print('deleting rule from {}: {}'.format(chain_name, rule))
    table = iptc.Table(iptc.Table.FILTER)
    with iptables_txn(table):
        chain = iptc.Chain(table, chain_name)
        for potential_rule in chain.rules:
            if Rule.from_iptc(potential_rule) == rule:
                chain.delete_rule(potential_rule)


def create_chain(chain_name):
    print('creating chain: {}'.format(chain_name))
    iptc.Table(iptc.Table.FILTER).create_chain(chain_name)


def delete_chain(chain_name):
    print('deleting chain: {}'.format(chain_name))
    chain = iptc.Chain(iptc.Table(iptc.Table.FILTER), chain_name)
    chain.flush()
    chain.delete()


def list_chain(chain_name):
    """List rules in a chain.

    Returns a list of iptables rules, or raises ChainDoesNotExist.
    """
    table = iptc.Table(iptc.Table.FILTER)
{'fraction_non_alphanumeric': 0.09915506958250497, 'fraction_numerical': 0.0, 'mean_word_length': 2.6871278057718735, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5685648', 'n_tokens_mistral': 2436, 'n_tokens_neox': 2352, 'n_words': 723}