#!/usr/bin/env python3

import errno
import smtplib
import socket
import time

import pytest


def send_test(smtp, testmsg):
    smtp.sendmail(
        'testsender@example.com',
        'testrcpt@example.com',
        testmsg.as_string(),
    )


def test_mode_normal(smtp, testmsg):
    smtp.ehlo()
    with pytest.raises(smtplib.SMTPRecipientsRefused):
        send_test(smtp, testmsg)


def test_mode_disabled(simta):
    with pytest.raises(smtplib.SMTPConnectError) as e:
        smtplib.SMTP('localhost', simta['port'])
    assert e.value.smtp_code == 554


def test_mode_global_relay(smtp, testmsg):
    send_test(smtp, testmsg)


def test_mode_tarpit(smtp, testmsg):
    startts = time.time()
    with pytest.raises(smtplib.SMTPDataError) as e:
        send_test(smtp, testmsg)
    assert e.value.smtp_code == 451
    assert time.time() - startts < 1


def test_mode_tarpit_timing(smtp, testmsg):
    startts = time.time()
    with pytest.raises(smtplib.SMTPDataError) as e:
        send_test(smtp, testmsg)
    assert e.value.smtp_code == 451
    assert time.time() - startts > 2.5
    assert time.time() - startts < 5


def test_mode_tempfail(smtp):
    smtp.ehlo()
    res = smtp.docmd('MAIL FROM:<>')
    assert res[0] == 451
    res = smtp.docmd('RCPT TO:<>')
    assert res[0] == 451
    res = smtp.docmd('DATA')
    assert res[0] == 503


def trigger_punishment(smtp):
    res = smtp.docmd('MAIL FROM:<testsender@example.com>')
    assert res[0] == 250
    res = smtp.docmd('RCPT TO:<badrcpt@example.edu>')
    assert res[0] == 551
    res = smtp.docmd('RCPT TO:<badrcpt@example.edu>')
    assert res[0] == 551


def test_punishment_mode_tempfail(smtp, testmsg):
    send_test(smtp, testmsg)
    trigger_punishment(smtp)
    test_mode_tempfail(smtp)


def test_punishment_trigger_mailfrom(smtp, testmsg):
    send_test(smtp, testmsg)
    res = smtp.docmd('MAIL FROM:<me@baddomain.example.com>')
    assert res[0] == 550
    res = smtp.docmd('MAIL FROM:<me@baddomain.example.com>')
    assert res[0] == 550
    test_mode_tempfail(smtp)


def test_punishment_trigger_auth(smtp, testmsg):
    send_test(smtp, testmsg)
    smtp.login('fakeuser', 'fakepassword')
    test_mode_tempfail(smtp)


# FIXME: This is very unreliable
@pytest.mark.xfail(reason='flaky', strict=False)
def test_punishment_trigger_nobanner(simta, testmsg):
    smtp = smtplib.SMTP('localhost', simta['port'])
    test_mode_global_relay(smtp, testmsg)
    smtp.quit()

    startts = time.time()
    conn = socket.create_connection(('localhost', simta['port']))
    conn.settimeout(5)
    conn.sendall(b'EHLO itsanevilclient\n')
    response = conn.recv(4096)
    assert response[:3] == b'220'
    response = conn.recv(4096)
    assert response[:3] == b'421'
    conn.sendall(b'MAIL FROM:<eviluser@example.com>\n')
    with pytest.raises(socket.error) as e:
        response = conn.recv(4096)
        response = conn.recv(4096)
    assert e.value.errno == errno.ECONNRESET
    assert time.time() - startts > 1
    assert time.time() - startts < 3


def test_punishment_mode_disabled(smtp_nocleanup, testmsg):
    smtp = smtp_nocleanup
    send_test(smtp, testmsg)
    trigger_punishment(smtp)
    res = smtp.docmd('RCPT TO:<badrcpt@example.edu>')
    assert res[0] == 421
    with pytest.raises(smtplib.SMTPServerDisconnected):
        res = smtp.docmd('RCPT TO:<badrcpt@example.edu>')


def test_punishment_mode_tarpit(smtp, testmsg):
    send_test(smtp, testmsg)
    trigger_punishment(smtp)
    smtp.docmd('RSET')
    test_mode_tarpit(smtp, testmsg)

@@ -2200,90 +2200,8 @@
 )%0A%0A%0A
-# FIXME: This is very unreliable%0A@pytest.mark.xfail(reason='flaky', strict=False)%0A
 def 
@@ -2542,32 +2542,45 @@
  conn.recv(4096)
+.splitlines()
 %0A    assert resp
@@ -2576,32 +2576,35 @@
 assert response%5B
+0%5D%5B
 :3%5D == b'220'%0A  
@@ -2609,34 +2609,211 @@
     
-response = conn.recv(4096)
+# Sometimes the first recv will get both the banner and the EHLO response.%0A    # Sometimes it won't.%0A    if len(response) == 1:%0A        response = conn.recv(4096)%0A    else:%0A        response = response%5B1%5D
 %0A   
@@ -2954,35 +2954,88 @@
     
-response = conn.recv(4096)%0A
+while True:%0A            conn.sendall(b'MAIL FROM:%3Ceviluser@example.com%3E%5Cn')%0A    
     
@@ -3090,19 +3090,20 @@
 e.errno 
-== 
+in %5B
 errno.EC
@@ -3110,16 +3110,30 @@
 ONNRESET
+, errno.EPIPE%5D
 %0A    ass

fix flaky test