Add tests for penaltybox and userthrottle

Author: flowerysong

penaltybox.c

@@ -25,6 +25,7 @@ int
 main(int ac, char *av[]) {
     char *      redis_host = "127.0.0.1";
     int         redis_port = 6379;
+    int         window = 300;
     char *      ip;
     char *      cksum;
     char *      from;
@@ -42,7 +43,7 @@ main(int ac, char *av[]) {
     urclHandle *urc;
     redisReply *res;
 
-    while ((c = getopt(ac, av, "h:H:p:P:")) != -1) {
+    while ((c = getopt(ac, av, "h:H:p:P:w:")) != -1) {
         switch (c) {
         case 'h':
             redis_host = optarg;
@@ -64,6 +65,13 @@ main(int ac, char *av[]) {
         case 'P':
             prefix = optarg;
             break;
+        case 'w':
+            errno = 0;
+            window = strtoll(optarg, NULL, 10);
+            if (errno != 0) {
+                err++;
+            }
+            break;
         case '?':
         default:
             err++;
@@ -130,7 +138,7 @@ main(int ac, char *av[]) {
         setenv("TZ", "", 1);
         tzset();
         then = mktime(&tm);
-        if ((timediff = difftime(now, then)) > 300) {
+        if ((timediff = difftime(now, then)) > window) {
             urcl_command(urc, key, "DEL %s", key);
             if (ht_threshold > 0) {
                 if (((res = urcl_command(urc, ht_key, "INCR %s", ht_key)) !=

test/test_penaltybox.py

@@ -0,0 +1,84 @@
+import subprocess
+import time
+
+import pytest
+
+
+@pytest.fixture
+def run_penaltybox(tool_path, redis):
+    def _run_penaltybox(cksum='foo', ip='10.11.112.113'):
+        env = {
+            'SIMTA_CHECKSUM': cksum,
+            'SIMTA_REMOTE_IP': ip,
+            'SIMTA_SMTP_MAIL_FROM': 'foo@example.com',
+        }
+
+        return subprocess.call(
+            [
+                tool_path('penaltybox'),
+                '-p', str(redis),
+                '-w', '1',
+                '-H', '3',
+                'no reason',
+            ],
+            env=env,
+        )
+
+    return _run_penaltybox
+
+
+def test_single(run_penaltybox):
+    assert run_penaltybox() == 1
+
+
+def test_window(run_penaltybox):
+    assert run_penaltybox() == 1
+    assert run_penaltybox() == 1
+
+
+@pytest.mark.parametrize('ip', [
+    '10.11.112.113',
+    'fdee:f057:3bf9:7dce:dead:60ff:d00d:60ff',
+])
+def test_accept(run_penaltybox, ip):
+    assert run_penaltybox(ip=ip) == 1
+    time.sleep(2)
+    assert run_penaltybox(ip=ip) == 0
+
+
+@pytest.mark.parametrize('ips', [
+    ('10.11.112.113', '10.11.112.1'),
+    ('fdee:f057:3bf9:7dce:dead:60ff:d00d:60ff', 'fdee:f057:3bf9:7dce:0001:0001:0001:0001'),
+])
+def test_accept_subnet(run_penaltybox, ips):
+    assert run_penaltybox(ip=ips[0]) == 1
+    time.sleep(2)
+    assert run_penaltybox(ip=ips[1]) == 0
+
+
+@pytest.mark.parametrize('ips', [
+    ('10.11.112.113', '10.11.111.1'),
+    ('fdee:f057:3bf9:7dce:dead:60ff:d00d:60ff', 'fdee:f057:3bf9:7dcd:0001:0001:0001:0001'),
+])
+def test_reject_subnet(run_penaltybox, ips):
+    assert run_penaltybox(ip=ips[0]) == 1
+    time.sleep(2)
+    assert run_penaltybox(ip=ips[1]) == 1
+
+
+def test_reject_checksum(run_penaltybox):
+    assert run_penaltybox() == 1
+    time.sleep(2)
+    assert run_penaltybox(cksum='bar') == 1
+
+
+def test_hattrick(run_penaltybox):
+    msgs = ['foo', 'bar', 'baz']
+    for msg in msgs:
+        assert run_penaltybox(cksum=msg) == 1
+    time.sleep(2)
+    for msg in msgs:
+        assert run_penaltybox(cksum=msg) == 0
+
+    # This IP should now be exempt from penalties
+    assert run_penaltybox(cksum='quux') == 0

test/test_userthrottle.py

@@ -0,0 +1,93 @@
+import subprocess
+
+import pytest
+
+
+@pytest.fixture
+def run_userthrottle(tool_path, redis):
+    def _run_userthrottle(auth='foo', ip='10.11.112.113', mfrom='foo@example.com', hfrom=None, nrcpts=1):
+        env = {
+            'SIMTA_AUTH_ID': auth,
+            'SIMTA_REMOTE_IP': ip,
+            'SIMTA_SMTP_MAIL_FROM': mfrom,
+            'SIMTA_NRCPTS': str(nrcpts),
+        }
+
+        if hfrom:
+            env['SIMTA_HEADER_FROM'] = hfrom
+
+        res = subprocess.check_output(
+            [
+                tool_path('userthrottle'),
+                '-p', str(redis),
+                '-d', 'example.com',
+            ],
+            env=env,
+        )
+
+        return res.strip()
+    return _run_userthrottle
+
+
+def test_single(run_userthrottle):
+    assert run_userthrottle() == '1'
+
+
+def test_multiple(run_userthrottle):
+    assert run_userthrottle() == '1'
+    assert run_userthrottle() == '2'
+
+
+def test_nrcpt(run_userthrottle):
+    assert run_userthrottle() == '1'
+    assert run_userthrottle(nrcpts=5) == '6'
+
+
+def test_multiple_ips(run_userthrottle):
+    assert run_userthrottle() == '1'
+    assert run_userthrottle(ip='10.12.112.114') == '2'
+
+
+def test_mfrom_mismatch(run_userthrottle):
+    assert run_userthrottle(mfrom='bar@example.com') == '2'
+
+
+def test_multiple_addresses(run_userthrottle):
+    assert run_userthrottle() == '1'
+    assert run_userthrottle(mfrom='bar@example.com') == '5'
+    assert run_userthrottle(mfrom='baz@example.com') == '17'
+    assert run_userthrottle() == '23'
+
+
+def test_multiple_auth(run_userthrottle):
+    assert run_userthrottle() == '1'
+    assert run_userthrottle(auth='bar', mfrom='bar@example.com') == '1'
+
+
+def test_multiple_mfrom(run_userthrottle):
+    assert run_userthrottle() == '1'
+    assert run_userthrottle(mfrom='bar@example.com') == '5'
+
+
+def test_hfrom_match(run_userthrottle):
+    assert run_userthrottle(hfrom='foo@example.com') == '1'
+    assert run_userthrottle(hfrom='foo@example.com') == '2'
+
+
+def test_hfrom_mismatch(run_userthrottle):
+    assert run_userthrottle(hfrom='bar@example.com') == '2'
+    assert run_userthrottle(hfrom='bar@example.com') == '4'
+
+
+def test_domain_mismatch(run_userthrottle):
+    assert run_userthrottle(mfrom='foo@example.edu') == '2'
+    assert run_userthrottle(hfrom='foo@example.com') == '4'
+
+
+def test_worst_case(run_userthrottle):
+    assert run_userthrottle() == '1'
+    assert run_userthrottle(mfrom='bar@example.edu') == '9'
+    assert run_userthrottle(mfrom='baz@example.edu', hfrom='baz@example.com') == '57'
+    assert run_userthrottle(mfrom='baz@example.edu', hfrom='baz@example.com', nrcpts=10) == '537'
+    assert run_userthrottle(auth='bar', mfrom='baz@example.edu', hfrom='baz@example.com') == '8'
+    assert run_userthrottle(auth='bar', mfrom='foo@example.edu', hfrom='baz@example.com') == '24'