diff --git a/.github/workflows/github-actions.yml b/.github/workflows/github-actions.yml new file mode 100644 index 0000000..5def6f6 --- /dev/null +++ b/.github/workflows/github-actions.yml @@ -0,0 +1,46 @@ +name: policyd-rate-limit +run-name: ${{ github.actor }} is running policyd-rate-limit CI tests +on: [push] +jobs: + flake8: + runs-on: ubuntu-latest + container: + image: python:bookworm + steps: + - uses: actions/checkout@v3 + - run: pip install tox + - run: apt-get update && apt-get install -y --no-install-recommends sudo + - run: useradd --uid 1000 testuser && mkdir -p /home/testuser && chown testuser -R . /home/testuser + - run: sudo -u testuser tox -e flake8 + check_rst: + runs-on: ubuntu-latest + container: + image: python:bookworm + steps: + - uses: actions/checkout@v3 + - run: pip install tox + - run: apt-get update && apt-get install -y --no-install-recommends sudo + - run: useradd --uid 1000 testuser && mkdir -p /home/testuser && chown testuser -R . /home/testuser + - run: sudo -u testuser tox -e check_rst + tests: + runs-on: ubuntu-latest + container: + image: python:bookworm + steps: + - uses: actions/checkout@v3 + - run: pip install tox + - run: apt-get update && apt-get install -y --no-install-recommends sudo + - run: useradd --uid 1000 testuser && mkdir -p /home/testuser && chown testuser -R . /home/testuser + - run: sudo -u testuser tox -e py3 + coverage: + runs-on: ubuntu-latest + container: + image: python:bookworm + steps: + - uses: actions/checkout@v3 + - run: pip install tox + - run: apt-get update && apt-get install -y --no-install-recommends sudo + - run: useradd --uid 1000 testuser && mkdir -p /home/testuser && chown testuser -R . /home/testuser + - run: sudo --preserve-env=COVERAGE_TOKEN -u testuser tox -e coverage + env: + COVERAGE_TOKEN: ${{ secrets.COVERAGE_TOKEN }} diff --git a/.travis.yml b/.travis.yml deleted file mode 100644 index 8af0563..0000000 --- a/.travis.yml +++ /dev/null @@ -1,25 +0,0 @@ -language: python -matrix: - include: - - python: "3.4" - env: TOX_ENV=flake8 - - python: "3.4" - env: TOX_ENV=check_rst - - python: "3.4" - env: TOX_ENV=py34 - - python: "3.5" - env: TOX_ENV=py35 - - python: "3.4" - env: TOX_ENV=coverage -cache: - directories: - - $HOME/.cache/pip/http/ - - $HOME/build/nitmir/policyd-rate-limit/.tox/$TOX_ENV/ -install: - - "travis_retry pip install setuptools --upgrade" - - "pip install tox" -script: - - tox -e $TOX_ENV -after_script: - - cat tox_log/*.log - diff --git a/Makefile b/Makefile index e2e4177..678537d 100644 --- a/Makefile +++ b/Makefile @@ -1,6 +1,11 @@ .PHONY: clean build install dist uninstall VERSION=`python3 setup.py -V` +WHL_FILES := $(wildcard dist/*.whl) +WHL_ASC := $(WHL_FILES:=.asc) +DIST_FILE := $(wildcard dist/*.tar.gz) +DIST_ASC := $(DIST_FILE:=.asc) + build: python3 setup.py build @@ -9,14 +14,19 @@ install: dist [ ! -f /etc/policyd-rate-limit.yaml ] && cp -n policyd_rate_limit/policyd-rate-limit.yaml /etc/ || true cp -n init/policyd-rate-limit /etc/init.d cp -n init/policyd-rate-limit.service /etc/systemd/system/ || true + cp -n init/policyd-rate-limit-clean.service /etc/systemd/system/policyd-rate-limit-clean.service + cp -n init/policyd-rate-limit-clean.timer /etc/systemd/system/policyd-rate-limit-clean.timer pip3 install policyd-rate-limit --no-cache-dir -U --force-reinstall --no-deps --no-binary :all -f ./dist/policyd-rate-limit-${VERSION}.tar.gz systemctl daemon-reload + systemctl enable policyd-rate-limit-clean.timer + systemctl start policyd-rate-limit-clean.timer uninstall: pip3 uninstall policyd-rate-limit || true reinstall: uninstall install purge: uninstall rm -f /etc/policyd-rate-limit.conf /etc/policyd-rate-limit.yaml rm -f /etc/init.d/policyd-rate-limit /etc/systemd/system/policyd-rate-limit.service + rm -f /etc/systemd/system/policyd-rate-limit-clean.service /etc/systemd/system/policyd-rate-limit-clean.timer rm -rf /var/lib/policyd-rate-limit/ clean_pyc: @@ -42,11 +52,8 @@ man_files: dist: python3 setup.py sdist -publish_pypi_release: - python setup.py sdist upload --sign - test_venv/bin/python: - virtualenv -p python3 test_venv + python3 -m venv test_venv test_venv/bin/pip3 install -U -r requirements-dev.txt test_venv: test_venv/bin/python @@ -56,3 +63,14 @@ coverage: clean_coverage test_venv export PATH=test_venv/bin/:$$PATH; echo $$PATH; pytest test_venv/bin/coverage html test_venv/bin/coverage report + +sign_release: $(WHL_ASC) $(DIST_ASC) + +dist/%.asc: + gpg --detach-sign -a $(@:.asc=) + +test_venv/bin/twine: test_venv + test_venv/bin/pip install twine + +publish_pypi_release: test_venv test_venv/bin/twine dist sign_release + test_venv/bin/twine upload --sign dist/* diff --git a/README.rst b/README.rst index c76311d..43ccfd9 100644 --- a/README.rst +++ b/README.rst @@ -1,7 +1,7 @@ Policyd rate limit ================== -|travis| |coverage| |github_version| |pypi_version| |license| +|coverage| |github_version| |pypi_version| |license| Postfix policyd server allowing to limit the number of mails accepted by postfix over several time periods, by sasl usernames and/or ip addresses. @@ -137,6 +137,14 @@ The ``.yaml`` are the new configuration format using the YAML syntax. * ``smtp_credentials``: Should we use credentials to connect to smtp_server ? if yes set ``["user", "password"]``, else ``null``. The default is ``null``. +* ``count_mode``: How sent mail are counted + + * ``0``: each RCPT TO are counted individualy. This is the how it was done historically. If set to 0, + the postfix check_policy_service must be set in smtpd_recipient_restrictions. + This is deprecated and should not be used anymore + * ``1``: recipient are counted in the DATA stage. The postfix parameter check_policy_service must be + defined in smtpd_data_restrictions. + This is the new default. Postfix settings ---------------- @@ -148,7 +156,7 @@ service. /etc/postfix/main.cf:: - smtpd_recipient_restrictions = + smtpd_data_restrictions = ..., check_policy_service { unix:ratelimit/policy, default_action=DUNNO }, ... @@ -158,15 +166,12 @@ On previous postfix versions, you must use: /etc/postfix/main.cf:: - smtpd_recipient_restrictions = + smtpd_data_restrictions = ..., check_policy_service unix:ratelimit/policy, ... -.. |travis| image:: https://badges.genua.fr/travis/nitmir/policyd-rate-limit/master.svg - :target: https://travis-ci.org/nitmir/policyd-rate-limit - .. |coverage| image:: https://badges.genua.fr/coverage/badge/policyd-rate-limit/master.svg :target: https://badges.genua.fr/coverage/policyd-rate-limit/ diff --git a/docs/policyd-rate-limit.8.rst b/docs/policyd-rate-limit.8.rst index d62d247..dabf47e 100644 --- a/docs/policyd-rate-limit.8.rst +++ b/docs/policyd-rate-limit.8.rst @@ -36,7 +36,7 @@ Setup For example, for postfix 3.0 and later, you can set in postfix **/etc/postfix/main.cf** configuration file:: - smtpd_recipient_restrictions = + smtpd_data_restrictions = ..., check_policy_service { unix:ratelimit/policy, default_action=DUNNO }, ... @@ -47,7 +47,7 @@ and will accept mail if policyd-rate-limit become unavailable. On previous postfix versions, you must use:: - smtpd_recipient_restrictions = + smtpd_data_restrictions = ..., check_policy_service unix:ratelimit/policy, ... diff --git a/docs/policyd-rate-limit.yaml.5.rst b/docs/policyd-rate-limit.yaml.5.rst index b76e1f4..a39a5d4 100644 --- a/docs/policyd-rate-limit.yaml.5.rst +++ b/docs/policyd-rate-limit.yaml.5.rst @@ -124,6 +124,14 @@ Settings if yes set ["user", "password"], else null. The default is null. +**count_mode** + How sent mail are counted. Set to **0**, each RCPT TO are counted individualy. + This is the how it was done historically. If set to 0, the postfix check_policy_service must be set in + smtpd_recipient_restrictions. This is deprecated and should not be used anymore. + Set to **1** recipient are counted in the DATA stage. The postfix parameter check_policy_service must be + defined in smtpd_data_restrictions. This is the new default. + + See also ======== diff --git a/init/policyd-rate-limit-clean.service b/init/policyd-rate-limit-clean.service new file mode 100644 index 0000000..dde73c4 --- /dev/null +++ b/init/policyd-rate-limit-clean.service @@ -0,0 +1,13 @@ +[Unit] +Description=Postfix policyd rate limiter - clean database +After=syslog.target + +[Service] +Type=oneshot +ExecStart=/usr/local/bin/policyd-rate-limit --clean +KillSignal=SIGINT +StandardOutput=syslog +StandardError=syslog + +[Install] +WantedBy=multi-user.target diff --git a/init/policyd-rate-limit-clean.timer b/init/policyd-rate-limit-clean.timer new file mode 100644 index 0000000..b29591c --- /dev/null +++ b/init/policyd-rate-limit-clean.timer @@ -0,0 +1,9 @@ +[Unit] +Description=Clean policyd rate limit database daily + +[Timer] +OnCalendar=daily +Persistent=true + +[Install] +WantedBy=timers.target diff --git a/policyd_rate_limit/config.py b/policyd_rate_limit/config.py index 64db4d6..974b79c 100644 --- a/policyd_rate_limit/config.py +++ b/policyd_rate_limit/config.py @@ -84,3 +84,9 @@ smtp_starttls = False # Should we use credentials to connect to smtp_server ? if yes set ("user", "password"), else None smtp_credentials = None + +# The time in seconds before an unused socket gets closed +delay_to_close = 300 + +# count mode. 0 for RCPT, 1 for DATA +count_mode = 0 diff --git a/policyd_rate_limit/policyd-rate-limit.yaml b/policyd_rate_limit/policyd-rate-limit.yaml index cf6e507..0d07dcc 100644 --- a/policyd_rate_limit/policyd-rate-limit.yaml +++ b/policyd_rate_limit/policyd-rate-limit.yaml @@ -109,3 +109,15 @@ smtp_server: ["localhost", 25] smtp_starttls: False # Should we use credentials to connect to smtp_server ? if yes set ["user", "password"], else null smtp_credentials: null + +# The time in seconds before an unused socket gets closed +delay_to_close: 300 + +# How sent mail are counted: +# * 0 each RCPT TO are counted individualy. This is the how it was done historically. If set to 0, +# the postfix check_policy_service must be set in smtpd_recipient_restrictions. +# This is deprecated and should not be used anymore +# * 1 recipient are counted in the DATA stage. The postfix parameter check_policy_service must be +# defined in smtpd_data_restrictions. +# This is the new default. +count_mode: 1 diff --git a/policyd_rate_limit/policyd.py b/policyd_rate_limit/policyd.py index 0595fd9..ee65987 100644 --- a/policyd_rate_limit/policyd.py +++ b/policyd_rate_limit/policyd.py @@ -24,10 +24,20 @@ class Pass(Exception): pass +class PolicydError(Exception): + pass + + +class PolicydConnectionClosed(PolicydError): + pass + + class Policyd(object): """The policy server class""" socket_data_read = {} socket_data_write = {} + last_used = {} + last_deprecation_warning = 0 def socket(self): """initialize the socket from the config parameters""" @@ -73,6 +83,19 @@ def close_connection(self, connection): pass connection.close() + def close_write_conn(self, connection): + """Removes a socket from the write dict""" + try: + del self.socket_data_write[connection] + except KeyError: + if config.debug: + sys.stderr.write( + ( + "Hmmm, a socket actually used to write a little " + "time ago wasn\'t in socket_data_write. Weird.\n" + ) + ) + def run(self): """The main server loop""" try: @@ -98,6 +121,9 @@ def run(self): sys.stderr.write('connection from %s\n' % (client_address,)) sys.stderr.flush() self.socket_data_read[connection] = [] + + # Updates the last_sed time for the socket. + self.last_used[connection] = time.time() # else there is data to read on a client socket else: self.read(socket) @@ -109,10 +135,24 @@ def run(self): if data_not_sent: self.socket_data_write[socket] = data_not_sent else: - self.close_connection(socket) + self.close_write_conn(socket) + + # Socket has been used, let's update its last_used time. + self.last_used[socket] = time.time() # the socket has been closed during read except KeyError: pass + # Closes unused socket for a long time. + __to_rm = [] + for (socket, last_used) in self.last_used.items(): + if socket == sock: + continue + if time.time() - last_used > config.delay_to_close: + self.close_connection(socket) + __to_rm.append(socket) + for socket in __to_rm: + self.last_used.pop(socket) + except (KeyboardInterrupt, utils.Exit): for socket in list(self.socket_data_read.keys()): if socket != self.sock: @@ -127,7 +167,7 @@ def read(self, connection): # read data data = connection.recv(1024).decode('UTF-8') if not data: - raise ValueError("connection closed") + raise PolicydConnectionClosed() if config.debug: sys.stderr.write(data) sys.stderr.flush() @@ -156,10 +196,17 @@ def read(self, connection): self.action(connection, request) else: self.socket_data_read[connection] = buffer + # Socket has been used, let's update its last_used time. + self.last_used[connection] = time.time() except (KeyboardInterrupt, utils.Exit): self.close_connection(connection) raise - except Exception as error: + except PolicydConnectionClosed: + if config.debug: + sys.stderr.write("Connection closed\n") + sys.stderr.flush() + self.close_connection(connection) + except Exception: traceback.print_exc() sys.stderr.flush() self.close_connection(connection) @@ -174,11 +221,50 @@ def action(self, connection, request): utils.database_init() with utils.cursor() as cur: try: - # only care if the protocol states is RCTP. If the policy delegation in postfix + # only care if the protocol states is RCTP or DATA. + # If the policy delegation in postfix # configuration is in smtpd_recipient_restrictions as said in the doc, # possible states are RCPT and VRFY. - if 'protocol_state' in request and request['protocol_state'].upper() != "RCPT": + # If in smtpd_data_restrictions only DATA is possible. + if 'protocol_state' not in request: + sys.stderr.write("Attribute 'protocol_state' not defined\n") + sys.stderr.flush() + raise Pass() + if config.count_mode not in {0, 1}: + sys.stderr.write("Settings 'count_mode' bad value %r\n" % ( + config.count_mode, + )) + sys.stderr.flush() + raise Pass() + if config.count_mode == 0 and request['protocol_state'].upper() != "RCPT": + if config.debug: + sys.stderr.write( + "Ignoring 'protocol_state' %r\n" % ( + request['protocol_state'].upper(), + ) + ) + sys.stderr.flush() raise Pass() + if config.count_mode == 1 and request['protocol_state'].upper() != "DATA": + if config.debug: + sys.stderr.write( + "Ignoring 'protocol_state' %r\n" % ( + request['protocol_state'].upper(), + ) + ) + sys.stderr.flush() + raise Pass() + if config.count_mode == 0: + if config.debug or time.time() - self.last_deprecation_warning > 60: + sys.stderr.write( + "WARNING: the 'count_mode' parameter is set to 0. " + "This is DEPRECATED. 'count_mode' should be set to 1 and postfix" + " config edited as stated in the README or " + "policyd-rate-limit.yaml(5)\n" + ) + sys.stderr.flush() + self.last_deprecation_warning = time.time() + # if user is authenticated, we filter by sasl username if config.limit_by_sasl and u'sasl_username' in request: id = request[u'sasl_username'] @@ -197,8 +283,12 @@ def action(self, connection, request): else: raise Pass() - #Custom limits per ID via MySQL + if request['protocol_state'].upper() == "RCPT": + recipient_count = 1 + elif request['protocol_state'].upper() == "DATA": + recipient_count = max(int(request["recipient_count"]), 1) + #Custom limits per ID via SQL custom_limits = config.limits_by_id if config.sql_limits_by_id != "": try: @@ -219,16 +309,20 @@ def action(self, connection, request): for mail_nb, delta in custom_limits.get(id, config.limits): cur.execute( ( - "SELECT COUNT(*) FROM mail_count " + "SELECT SUM(recipient_count) FROM mail_count " "WHERE id = %s AND date >= %s" ) % ((config.format_str,)*2), (id, int(time.time() - delta)) ) - nb = cur.fetchone()[0] + nb = cur.fetchone()[0] or 0 if config.debug: - sys.stderr.write("%03d/%03d hit since %ss\n" % (nb, mail_nb, delta)) + sys.stderr.write( + "%03d/%03d hit since %ss\n" % ( + nb + recipient_count, mail_nb, delta + ) + ) sys.stderr.flush() - if nb >= mail_nb: + if nb + recipient_count > mail_nb: action = config.fail_action if config.report and delta in config.report_limits: utils.hit(cur, delta, id) @@ -241,8 +335,28 @@ def action(self, connection, request): sys.stderr.write(u"insert id %s\n" % id) sys.stderr.flush() cur.execute( - "INSERT INTO mail_count VALUES (%s, %s)" % ((config.format_str,)*2), - (id, int(time.time())) + "INSERT INTO mail_count VALUES (%s, %s, %s, %s, %s)" % ( + (config.format_str,)*5 + ), + ( + id, int(time.time()), recipient_count, + request.get("instance", "empty"), request['protocol_state'] + ) + ) + # If action is a failure and using legacy mode, remove previous + # recorded event for this mail in the + # database. The mail has not been sent, we should not count any recipient + if ( + config.count_mode == 0 and + action == config.fail_action and + request['protocol_state'].upper() == "RCPT" and + request.get("instance") + ): + cur.execute( + "DELETE FROM mail_count WHERE instance = %s AND protocol_state = %s" % ( + (config.format_str,)*2 + ), + (request["instance"], request['protocol_state']) ) except utils.cursor.backend_module.Error as error: utils.cursor.del_db() @@ -255,3 +369,7 @@ def action(self, connection, request): # return the result to the client self.socket_data_write[connection] = data.encode('UTF-8') + # Wipe the read buffer (otherwise it'll be added up for eternity) + self.socket_data_read[connection].clear() + # Socket has been used, let's update its last_used time. + self.last_used[connection] = time.time() diff --git a/policyd_rate_limit/tests/test_daemon.py b/policyd_rate_limit/tests/test_daemon.py index 095c2da..155b39f 100644 --- a/policyd_rate_limit/tests/test_daemon.py +++ b/policyd_rate_limit/tests/test_daemon.py @@ -30,6 +30,7 @@ def setUp(self): report_limits=[60, 86400], user="root", group="root", + count_mode=0, ) def tearDown(self): @@ -41,14 +42,15 @@ def test_main_unix_socket(self): self.base_test(cfg) def test_main_afinet_socket(self): - self.base_config["SOCKET"] = ("127.0.0.1", 27184) + self.base_config["SOCKET"] = ["127.0.0.1", 27184] with test_utils.lauch(self.base_config) as cfg: self.base_test(cfg) - def test_main_afinet6_socket(self): - self.base_config["SOCKET"] = ("::1", 27184) - with test_utils.lauch(self.base_config) as cfg: - self.base_test(cfg) + # travis CI/Github Action has no IPv6 support + # def test_main_afinet6_socket(self): + # self.base_config["SOCKET"] = ["::1", 27184] + # with test_utils.lauch(self.base_config) as cfg: + # self.base_test(cfg) def test_no_debug_no_report(self): self.base_config["debug"] = False @@ -56,6 +58,55 @@ def test_no_debug_no_report(self): with test_utils.lauch(self.base_config) as cfg: self.base_test(cfg) + def test_limit(self): + with test_utils.lauch(self.base_config) as cfg: + for i in range(10): + data = test_utils.send_policyd_request(cfg["SOCKET"], sasl_username="test") + self.assertEqual(data.strip(), b"action=dunno") + # the eleventh counted requests should fail + data = test_utils.send_policyd_request(cfg["SOCKET"], sasl_username="test") + self.assertEqual(data.strip(), b"action=defer_if_permit Rate limit reach, retry later") + + def test_limit_batch(self): + with test_utils.lauch(self.base_config) as cfg: + # Send a batch of mails + for i in range(10): + data = test_utils.send_policyd_request( + cfg["SOCKET"], sasl_username="test", instance="test" + ) + self.assertEqual(data.strip(), b"action=dunno") + # the eleventh counted requests should fail and the 10 previous should be discard + data = test_utils.send_policyd_request( + cfg["SOCKET"], sasl_username="test", instance="test" + ) + self.assertEqual(data.strip(), b"action=defer_if_permit Rate limit reach, retry later") + # The limit should have be reverted (cf instance) + for i in range(10): + data = test_utils.send_policyd_request(cfg["SOCKET"], sasl_username="test") + self.assertEqual(data.strip(), b"action=dunno") + # the eleventh counted requests should fail + data = test_utils.send_policyd_request(cfg["SOCKET"], sasl_username="test") + self.assertEqual(data.strip(), b"action=defer_if_permit Rate limit reach, retry later") + + def test_limit_batch2(self): + self.base_config["count_mode"] = 1 + with test_utils.lauch(self.base_config) as cfg: + # Send a batch of mails + data = test_utils.send_policyd_request( + cfg["SOCKET"], sasl_username="test", protocol_state="DATA", recipient_count=11 + ) + self.assertEqual(data.strip(), b"action=defer_if_permit Rate limit reach, retry later") + # The limit should have be reverted (cf instance) + data = test_utils.send_policyd_request( + cfg["SOCKET"], sasl_username="test", protocol_state="DATA", recipient_count=10 + ) + self.assertEqual(data.strip(), b"action=dunno") + # the eleventh counted requests should fail + data = test_utils.send_policyd_request( + cfg["SOCKET"], sasl_username="test", protocol_state="DATA", recipient_count=1 + ) + self.assertEqual(data.strip(), b"action=defer_if_permit Rate limit reach, retry later") + def test_slow_connection(self): with test_utils.lauch(self.base_config) as cfg: with test_utils.sock(cfg["SOCKET"]) as s: @@ -166,7 +217,7 @@ def test_already_running(self): f.write("") os.chmod(self.base_config["pidfile"], 0) with test_utils.lauch(self.base_config, get_process=True) as p: - self.assertEqual(p.wait(), 6) + self.assertEqual(p.wait(timeout=5), 6) finally: try: os.remove(self.base_config["pidfile"]) @@ -174,10 +225,10 @@ def test_already_running(self): pass def test_bad_socket_bind_address(self): - self.base_config["SOCKET"] = ("toto", 1234) + self.base_config["SOCKET"] = ["toto", 1234] with test_utils.lauch(self.base_config, get_process=True, no_wait=True) as p: self.assertEqual(p.wait(), 4) - self.base_config["SOCKET"] = ("192.168::1", 1234) + self.base_config["SOCKET"] = ["192.168::1", 1234] with test_utils.lauch(self.base_config, get_process=True, no_wait=True) as p: self.assertEqual(p.wait(), 6) diff --git a/policyd_rate_limit/tests/utils.py b/policyd_rate_limit/tests/utils.py index d6c1879..56e0658 100644 --- a/policyd_rate_limit/tests/utils.py +++ b/policyd_rate_limit/tests/utils.py @@ -16,6 +16,8 @@ import tempfile import subprocess import time +import string +import random from contextlib import contextmanager @@ -28,9 +30,9 @@ helo_name=mail.example.com sender=bar@example.com recipient=foo@example.com -recipient_count=0 +recipient_count=%(recipient_count)s queue_id= -instance=fd3.57cea9c4.143ea.0 +instance=%(instance)s size=0 etrn_domain= stress= @@ -48,16 +50,31 @@ """ -def postfix_request(sasl_username="", client_address="127.0.0.1", protocol_state="RCPT"): +def postfix_request( + sasl_username="", client_address="127.0.0.1", protocol_state="RCPT", + instance=None, recipient_count=None, +): + if instance is None: + letters = string.ascii_letters + string.digits + '.' + instance = ''.join(random.choice(letters) for _ in range(16)) + if recipient_count is None: + if protocol_state == "DATA": + recipient_count = 1 + else: + recipient_count = 0 return (POSTFIX_TEMPLATE % { "sasl_username": sasl_username, "client_address": client_address, - "protocol_state": protocol_state + "protocol_state": protocol_state, + "instance": instance, + "recipient_count": recipient_count, }).encode("utf-8") @contextmanager def sock(addr): + if isinstance(addr, list): + addr = tuple(addr) if isinstance(addr, str): s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) elif '.' in addr[0]: @@ -73,10 +90,15 @@ def sock(addr): s.close() -def send_policyd_request(addr, sasl_username="", client_address="127.0.0.1", protocol_state="RCPT"): +def send_policyd_request( + addr, sasl_username="", client_address="127.0.0.1", protocol_state="RCPT", + instance=None, recipient_count=None, +): with sock(addr) as s: s.send( - postfix_request(sasl_username, client_address, protocol_state) + postfix_request( + sasl_username, client_address, protocol_state, instance, recipient_count + ) ) data = s.recv(1024) return data @@ -92,7 +114,7 @@ def gen_config(new_config): os.path.join(os.path.dirname(__file__), '..', 'policyd-rate-limit.yaml') ) with open(default_config) as f: - config = yaml.load(f) + config = yaml.load(f, Loader=yaml.SafeLoader) config.update(new_config) cfg_path = tempfile.mktemp('.yaml') with open(cfg_path, 'w') as f: @@ -156,7 +178,7 @@ def lauch(new_config, get_process=False, options=None, no_coverage=False, no_wai try: if cfg_path: with open(cfg_path) as f: - cfg = yaml.load(f) + cfg = yaml.load(f, Loader=yaml.SafeLoader) if not no_wait: time.sleep(0.01) for i in range(100): diff --git a/policyd_rate_limit/utils.py b/policyd_rate_limit/utils.py index 571fe51..044771c 100644 --- a/policyd_rate_limit/utils.py +++ b/policyd_rate_limit/utils.py @@ -85,7 +85,7 @@ def __init__(self, config_file=None): # new config file use yaml else: with open(config_file) as f: - self._config = yaml.load(f) + self._config = yaml.load(f, Loader=yaml.SafeLoader) self.config_file = config_file break except PermissionError: @@ -268,7 +268,7 @@ def set_db(cls, value): def del_db(cls): try: cls._db[threading.current_thread()].close() - except: + except Exception: pass try: del cls._db[threading.current_thread()] @@ -327,14 +327,18 @@ def clean(): max_delta = max(max_delta, delta) # remove old record older than 2*max_delta expired = int(time.time() - max_delta - max_delta) + report_text = "" with cursor() as cur: cur.execute("DELETE FROM mail_count WHERE date <= %s" % config.format_str, (expired,)) print("%d records deleted" % cur.rowcount) # if report is True, generate a mail report if config.report and config.report_to: - send_report(cur) - # The mail report has been successfully send, flush limit_report - cur.execute("DELETE FROM limit_report") + report_text = gen_report(cur) + # The mail report has been successfully send, flush limit_report + cur.execute("DELETE FROM limit_report") + # send report + if len(report_text) != 0: + send_report(report_text) try: if config.backend == PGSQL_DB: @@ -355,10 +359,11 @@ def clean(): cursor.get_db().autocommit = False -def send_report(cur): +def gen_report(cur): cur.execute("SELECT id, delta, hit FROM limit_report") # list to sort ids by hits report = list(cur.fetchall()) + text = [] if not config.report_only_if_needed or report: if report: text = ["Below is the table of users who hit a limit since the last cleanup:", ""] @@ -404,46 +409,50 @@ def send_report(cur): else: text = ["No user hit a limit since the last cleanup"] text.extend(["", "-- ", "policyd-rate-limit"]) + return text - # check that smtp_server is wekk formated - if isinstance(config.smtp_server, (list, tuple)): - if len(config.smtp_server) >= 2: - server = smtplib.SMTP(config.smtp_server[0], config.smtp_server[1]) - elif len(config.smtp_server) == 1: - server = smtplib.SMTP(config.smtp_server[0], 25) - else: - raise ValueError("bad smtp_server should be a tuple (server_adress, port)") + +def send_report(text): + # check that smtp_server is wekk formated + if isinstance(config.smtp_server, (list, tuple)): + if len(config.smtp_server) >= 2: + server = smtplib.SMTP(config.smtp_server[0], config.smtp_server[1]) + elif len(config.smtp_server) == 1: + server = smtplib.SMTP(config.smtp_server[0], 25) else: raise ValueError("bad smtp_server should be a tuple (server_adress, port)") + else: + raise ValueError("bad smtp_server should be a tuple (server_adress, port)") - try: - # should we use starttls ? - if config.smtp_starttls: - server.starttls() - # should we use credentials ? - if config.smtp_credentials: - if ( - isinstance(config.smtp_credentials, (list, tuple)) and - len(config.smtp_credentials) >= 2 - ): - server.login(config.smtp_credentials[0], config.smtp_credentials[1]) - else: - ValueError("bad smtp_credentials should be a tuple (login, password)") - - if not isinstance(config.report_to, list): - report_to = [config.report_to] + try: + # should we use starttls ? + if config.smtp_starttls: + server.starttls() + # should we use credentials ? + if config.smtp_credentials: + if ( + isinstance(config.smtp_credentials, (list, tuple)) and + len(config.smtp_credentials) >= 2 + ): + server.login(config.smtp_credentials[0], config.smtp_credentials[1]) else: - report_to = config.report_to - for rcpt in report_to: - # Start building the mail report - msg = MIMEMultipart() - msg['Subject'] = config.report_subject or "" - msg['From'] = config.report_from or "" - msg['To'] = rcpt - msg.attach(MIMEText("\n".join(text), 'plain')) - server.sendmail(config.report_from or "", rcpt, msg.as_string()) - finally: - server.quit() + ValueError("bad smtp_credentials should be a tuple (login, password)") + + if not isinstance(config.report_to, list): + report_to = [config.report_to] + else: + report_to = config.report_to + for rcpt in report_to: + # Start building the mail report + msg = MIMEMultipart() + msg['Subject'] = config.report_subject or "" + msg['From'] = config.report_from or "" + msg['To'] = rcpt + msg.attach(MIMEText("\n".join(text), 'plain')) + server.sendmail(config.report_from or "", rcpt, msg.as_string()) + finally: + print('report is sent') + server.quit() def database_init(): @@ -451,30 +460,64 @@ def database_init(): with cursor() as cur: query = """CREATE TABLE IF NOT EXISTS mail_count ( id varchar(40) NOT NULL, - date bigint NOT NULL + date bigint NOT NULL, + recipient_count int DEFAULT 1, + instance varchar(40) NOT NULL, + protocol_state varchar(10) NOT NULL );""" - if config.sql_limits_by_id != "": - if config.backend == MYSQL_DB: - query_limits = """CREATE TABLE IF NOT EXISTS rate_limits ( - id int NOT NULL AUTO_INCREMENT, - limits varchar(255) NOT NULL, - PRIMARY KEY (id) - );""" - if config.sql_limits_by_id != "": - if config.backend == PGSQL_DB: - query_limits = """CREATE TABLE IF NOT EXISTS rate_limits ( - id int NOT NULL, - limits varchar(255) NOT NULL, - PRIMARY KEY (id) - );""" + if config.backend == MYSQL_DB: + query_limits = """CREATE TABLE IF NOT EXISTS rate_limits ( + id int NOT NULL AUTO_INCREMENT, + limits varchar(255) NOT NULL, + PRIMARY KEY (id) + );""" + else: + query_limits = """CREATE TABLE IF NOT EXISTS rate_limits ( + id int NOT NULL, + limits varchar(255) NOT NULL, + PRIMARY KEY (id) + );""" # if report is enable, also create the table for storing report datas - if config.report: - query_report = """CREATE TABLE IF NOT EXISTS limit_report ( - id varchar(40) NOT NULL, - delta int NOT NULL, - hit int NOT NULL DEFAULT 0 - );""" + query_report = """CREATE TABLE IF NOT EXISTS limit_report ( + id varchar(40) NOT NULL, + delta int NOT NULL, + hit int NOT NULL DEFAULT 0 + );""" + # Test the table version + try: + cur.execute("SELECT recipient_count FROM mail_count") + except cursor.backend_module.Error as error: + # If the table mail_count exists but the new column + # recipient_count does not, drop the table (it only + # contains temporary data). It will be recreated below. + if ( + (cursor.backend == MYSQL_DB and error.args[0] == 1054) or + ( + cursor.backend == SQLITE_DB and + error.args[0] == 'no such column: recipient_count' + ) or + ( + cursor.backend == PGSQL_DB and + isinstance(error, cursor.backend_module.errors.UndefinedColumn) + ) + + ): + cursor.get_db().commit() + cur.execute("DROP TABLE mail_count") + # Mysql table 'mail_count' doesn't exist + elif cursor.backend == MYSQL_DB and error.args[0] == 1146: + cursor.get_db().commit() + elif cursor.backend == SQLITE_DB and error.args[0] == 'no such table: mail_count': + cursor.get_db().commit() + elif ( + cursor.backend == PGSQL_DB and + isinstance(error, cursor.backend_module.errors.UndefinedTable) + ): + cursor.get_db().commit() + else: + raise + # Create the table if needed try: if cursor.backend == MYSQL_DB: # ignore possible warnings about the table already existing diff --git a/requirements-dev.txt b/requirements-dev.txt index c3fc055..fa275f7 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,6 +1,5 @@ pyyaml pytest pytest-pythonpath -pytest-warnings coverage diff --git a/setup.py b/setup.py index 9ece519..dff2e1d 100755 --- a/setup.py +++ b/setup.py @@ -37,7 +37,7 @@ def add_data_file(dir, file, check_dir=False, mkdir=False): setup( name='policyd-rate-limit', - version='0.7.1', + version='1.1.0', description=DESC, long_description=README, author='Valentin Samir', diff --git a/tox.ini b/tox.ini index 66cbdb4..09c56af 100644 --- a/tox.ini +++ b/tox.ini @@ -2,32 +2,21 @@ envlist= flake8, check_rst, - py34, - py35 + py3 + [flake8] max-line-length=100 + [base] deps = -r{toxinidir}/requirements-dev.txt -[post_cmd] -commands= - find {toxworkdir} -name '*.pyc' -delete - find {toxworkdir} -name __pycache__ -delete - mkdir -p {toxinidir}/tox_logs/ - bash -c "mv {toxworkdir}/{envname}/log/* {toxinidir}/tox_logs/" -whitelist_externals= - find - bash - mkdir [testenv] commands= py.test -rw {posargs:policyd_rate_limit/tests/} coverage report - {[post_cmd]commands} -whitelist_externals={[post_cmd]whitelist_externals} [testenv:flake8] @@ -36,8 +25,6 @@ deps=flake8 skip_install=True commands= flake8 {toxinidir}/policyd_rate_limit {toxinidir}/policyd-rate-limit - {[post_cmd]commands} -whitelist_externals={[post_cmd]whitelist_externals} [testenv:check_rst] @@ -48,16 +35,14 @@ deps= skip_install=True commands= rst2html.py --strict {toxinidir}/README.rst /dev/null - {[post_cmd]commands} -whitelist_externals={[post_cmd]whitelist_externals} -[testenv:py34] -basepython=python3.4 +[testenv:py3] +basepython=python3 deps = {[base]deps} -[testenv:py35] -basepython=python3.5 +[testenv:py39] +basepython=python3.9 deps = {[base]deps} @@ -75,5 +60,5 @@ commands= coverage report coverage html {toxinidir}/.update_coverage "{toxinidir}" "policyd-rate-limit" - {[post_cmd]commands} -whitelist_externals={[post_cmd]whitelist_externals} +allowlist_externals= + {toxinidir}/.update_coverage