mod common;

pub use common::*;

use byte_strings::c_str;
use indymilter_test::*;
use spf_milter::*;
use std::net::{Ipv4Addr, Ipv6Addr};
use viaspf::lookup::{LookupError, Name};

#[tokio::test]
async fn verify_sender() {
    let opts = configure_logging(CliOptions::builder())
        .config_file(to_config_file_name(file!()))
        .build();

    let lookup = MockLookup::new()
        .with_lookup_txt(|name| match name.as_str() {
            "amy.org." | "mallory.org." => Ok(vec!["v=spf1 mx -all".into()]),
            "mail.amy.org." | "mail.mallory.org." => Ok(vec!["v=spf1 a -all".into()]),
            "error.mallory.org." => Ok(vec!["v=spf1 a -all".into(), "v=spf1 mx -all".into()]),
            _ => Err(LookupError::NoRecords),
        })
        .with_lookup_mx(|name| match name.as_str() {
            "amy.org." => Ok(vec![Name::new("mail.amy.org").unwrap()]),
            "mallory.org." => Ok(vec![Name::new("mail.mallory.org").unwrap()]),
            _ => Err(LookupError::NoRecords),
        })
        .with_lookup_a(|name| match name.as_str() {
            "mail.amy.org." => Ok(vec![Ipv4Addr::new(123, 123, 123, 123)]),
            _ => Err(LookupError::NoRecords),
        })
        .with_lookup_aaaa(|name| match name.as_str() {
            "mail.mallory.org." => Ok(vec![Ipv6Addr::new(0x123, 0, 0, 0, 0, 0, 0x123, 0x123)]),
            _ => Err(LookupError::NoRecords),
        });

    let config = Config::read_with_lookup(opts, lookup).await.unwrap();

    let milter = SpfMilter::spawn(config).await.unwrap();

    let addr = milter.addr();

    // Test 1: Accept authorised sender (HELO)

    let mut conn = TestConnection::open(addr).await.unwrap();

    conn.macros(MacroStage::Connect, [("j", "mail.gluet.ch")]).await.unwrap();

    let status = conn.connect("client.amy.org", [123, 123, 123, 123]).await.unwrap();
    assert_eq!(status, Status::Continue);

    let status = conn.helo("mail.amy.org").await.unwrap();
    assert_eq!(status, Status::Continue);

    // For completeness, test repeated use of SMTP `HELO` (see debug log):
    let status = conn.helo("mail.amy.org").await.unwrap();
    assert_eq!(status, Status::Continue);

    let status = conn.mail(["<from@amy.org>"]).await.unwrap();
    assert_eq!(status, Status::Continue);

    conn.macros(MacroStage::Data, [("i", "1234567ABC")]).await.unwrap();

    let (actions, status) = conn.eom().await.unwrap();
    assert_eq!(status, Status::Continue);

    assert!(actions.has_insert_header(
        0,
        "Authentication-Results",
        "mail.gluet.ch; spf=pass smtp.helo=mail.amy.org"
    ));

    conn.close().await.unwrap();

    // Test 2: Accept authorised sender (MAIL FROM)

    let mut conn = TestConnection::open(addr).await.unwrap();

    conn.macros(MacroStage::Connect, [("j", "mail.gluet.ch")]).await.unwrap();

    let status = conn.connect("client.amy.org", [123, 123, 123, 123]).await.unwrap();
    assert_eq!(status, Status::Continue);

    let status = conn.helo("mail.example.com").await.unwrap();
    assert_eq!(status, Status::Continue);

    let status = conn.mail(["<from@amy.org>"]).await.unwrap();
    assert_eq!(status, Status::Continue);

    conn.macros(MacroStage::Data, [("i", "1234567ABC")]).await.unwrap();

    let (actions, status) = conn.eom().await.unwrap();
    assert_eq!(status, Status::Continue);

    assert!(actions.has_insert_header(
        0,
        "Authentication-Results",
        "mail.gluet.ch; spf=pass smtp.mailfrom=amy.org"
    ));

    conn.close().await.unwrap();

    // Test 3: Reject unauthorised sender (HELO)

    let mut conn = TestConnection::open(addr).await.unwrap();

    conn.macros(MacroStage::Connect, [("j", "mail.gluet.ch")]).await.unwrap();

    let status =
        conn.connect("client.mallory.org", [0x123, 0, 0, 0, 0, 0, 0x123, 0x456]).await.unwrap();
    assert_eq!(status, Status::Continue);

    let status = conn.helo("mail.mallory.org").await.unwrap();
    assert_eq!(
        status,
        Status::Reject {
            message: Some(c_str!("550 5.7.23 SPF validation failed").into()),
        }
    );

    conn.close().await.unwrap();

    // Test 4: Reject unauthorised sender (MAIL FROM)

    let mut conn = TestConnection::open(addr).await.unwrap();

    conn.macros(MacroStage::Connect, [("j", "mail.gluet.ch")]).await.unwrap();

    let status =
        conn.connect("client.mallory.org", [0x123, 0, 0, 0, 0, 0, 0x123, 0x456]).await.unwrap();
    assert_eq!(status, Status::Continue);

    let status = conn.helo("mail.example.com").await.unwrap();
    assert_eq!(status, Status::Continue);

    let status = conn.mail(["<from@mallory.org>"]).await.unwrap();
    assert_eq!(
        status,
        Status::Reject {
            message: Some(c_str!("550 5.7.23 SPF validation failed").into()),
        }
    );

    conn.close().await.unwrap();

    // Test 5: Accept exempt unauthorised sender (MAIL FROM)

    let mut conn = TestConnection::open(addr).await.unwrap();

    conn.macros(MacroStage::Connect, [("j", "mail.gluet.ch")]).await.unwrap();

    let status =
        conn.connect("client.mallory.org", [0x123, 0, 0, 0, 0, 0, 0x123, 0x456]).await.unwrap();
    assert_eq!(status, Status::Continue);

    let status = conn.helo("mail.example.com").await.unwrap();
    assert_eq!(status, Status::Continue);

    // Accept message from unauthorised but exempt sender.
    let status = conn.mail(["<postmaster@mail.mallory.org>"]).await.unwrap();
    assert_eq!(status, Status::Continue);

    conn.macros(MacroStage::Data, [("i", "1234567ABC")]).await.unwrap();

    let (actions, status) = conn.eom().await.unwrap();
    assert_eq!(status, Status::Continue);

    assert!(!actions.has_insert_header(any(), "Authentication-Results", any()));

    conn.close().await.unwrap();

    // Test 6: Reject misconfigured sender (MAIL FROM)

    let mut conn = TestConnection::open(addr).await.unwrap();

    conn.macros(MacroStage::Connect, [("j", "mail.gluet.ch")]).await.unwrap();

    let status =
        conn.connect("client.mallory.org", [0x123, 0, 0, 0, 0, 0, 0x123, 0x456]).await.unwrap();
    assert_eq!(status, Status::Continue);

    let status = conn.helo("mail.example.com").await.unwrap();
    assert_eq!(status, Status::Continue);

    let status = conn.mail(["<from@error.mallory.org>"]).await.unwrap();
    assert_eq!(
        status,
        Status::Reject {
            message: Some(
                c_str!("550 5.7.24 SPF validation error: more than one SPF record found").into()
            ),
        }
    );

    conn.close().await.unwrap();

    // Test 7: Multiple messages per connection

    let mut conn = TestConnection::open(addr).await.unwrap();

    conn.macros(MacroStage::Connect, [("j", "mail.gluet.ch")]).await.unwrap();

    let status = conn.connect("client.amy.org", [123, 123, 123, 123]).await.unwrap();
    assert_eq!(status, Status::Continue);

    let status = conn.helo("mail.example.com").await.unwrap();
    assert_eq!(status, Status::Continue);

    // First message from an authorised sender is accepted.
    let status = conn.mail(["<from@amy.org>"]).await.unwrap();
    assert_eq!(status, Status::Continue);

    conn.macros(MacroStage::Data, [("i", "1234567ABC")]).await.unwrap();

    let (actions, status) = conn.eom().await.unwrap();
    assert_eq!(status, Status::Continue);

    assert!(actions.has_insert_header(
        0,
        "Authentication-Results",
        "mail.gluet.ch; spf=pass smtp.mailfrom=amy.org"
    ));

    // Second message from an authorised sender is aborted.
    let status = conn.mail(["<from@amy.org>"]).await.unwrap();
    assert_eq!(status, Status::Continue);

    conn.abort().await.unwrap();

    // Third message from an unauthorised sender is rejected.
    let status = conn.mail(["<from@mallory.org>"]).await.unwrap();
    assert_eq!(
        status,
        Status::Reject {
            message: Some(c_str!("550 5.7.23 SPF validation failed").into()),
        }
    );

    conn.close().await.unwrap();

    milter.shutdown().await.unwrap();
}