#!/usr/bin/env python3

import milter.client

# milter.client.Sessionを継承したクラスを実装する。
class MilterTemplate(milter.client.Session):
    # contextの後の引数は、milter.client.Clientのregisterメソッドで渡す
    # 2番目以降の引数に対応する。
    # (client.register(MilterTemplate, options))
    def __init__(self, context, options):
        super().__init__(context)

    # 1つのSMTPセッションの中で複数メールが送信される場合に、その切れ目で呼ばれる。
    # 初期化(super().__init__())時にも呼ばれる。
    # 1メール毎にセッションを分けるか、1セッションの中で複数メールを送るかは
    # クライアント次第。
    def reset(self):
        pass

    # 各処理でのエラー発生時に呼ばれる。
    def on_error(self, event, exception):
        pass

    # 以下の関数は、SMTPプロトコルの各段階と対応している。

    # MAIL FROM
    def envelope_from(self, from_):
        pass

    # RCPT TO
    def envelope_recipient(self, to):
        pass

    # DATA
    # MAIL FROMとRCPT TOが揃った時点であるため、初期化処理を行うのに良いタイミング。
    def data(self):
        pass

    # ヘッダー1行毎に複数回呼ばれる。
    def header(self, name, value):
        pass

    # ヘッダーを全て処理した時点で呼ばれる。
    def end_of_header(self):
        pass

    # 本文。ある程度大きい場合は、チャンクに分かれて複数回呼ばれる。
    def body(self, chunk):
        pass

    # 本文を全て処理した時点で呼ばれる。
    def end_of_message(self):
        pass


# コマンドラインをパースして、milter.client.Clientインスタンスを生成する。
# Clientインスタンスは、SMTPのセッション毎に
# milter.client.Sessionインスタンスを立てて処理を行わせる。
command_line = milter.client.CommandLine()
with command_line.run() as (client, options):
    # このClientインスタンスが使うSessionクラスを登録する。
    # SMTPのセッションが開始されると、ここで登録したクラスを
    # コンストラクトして処理をさせることになる。
    # そのためmilter開発の基本的な流れは、
    # milter.client.Sessionを継承したクラスを実装して、ここで登録することになる。
    client.register(MilterTemplate, options)

#!/usr/bin/env python3
#
# Copyright (C) 2022  Sutou Kouhei <kou@clear-code.com>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import re

import milter.client

class MilterReplace(milter.client.Session):
    def __init__(self, context, patterns):
        super().__init__(context)
        self._patterns = patterns

    def header(self, name, value):
        self._headers.append([name, value])

    def body(self, chunk):
        self._body += chunk

    def end_of_message(self):
        header_indexes = {}
        for name, value in self._headers:
            if name not in header_indexes:
                header_indexes[name] = 0
            header_indexes[name] += 1
            for pattern, replaced in self._patterns.items():
                replaced_value, _ = pattern.subn(replaced, value)
                if value != replaced_value:
                    self._change_header(name,
                                        header_indexes[name],
                                        replaced_value)
                    break

        for pattern, replaced in self._patterns.items():
            body = self._body.decode("utf-8")
            replaced_body, _ = pattern.subn(replaced, body)
            if body != replaced_body:
                self._replace_body(replaced_body)

    def reset(self):
        self._headers = []
        self._body = b""

command_line = milter.client.CommandLine()
with command_line.run() as (client, options):
    # "world" を "ClearCode" に置換する！
    patterns = {
        re.compile("world", re.IGNORECASE): "ClearCode",
    }
    client.register(MilterReplace, patterns)

def end_of_message(self):
    self._change_header("Subject", 1, "new subject")

# 1つのSMTPセッションの中で複数メールが送信される場合に、その切れ目で呼ばれる。
# 初期化(super().__init__())時にも呼ばれる。
# 1メール毎にセッションを分けるか、1セッションの中で複数メールを送るかは
# クライアント次第。
def reset(self):
    # ヘッダーを管理するためのメンバー変数を初期化する。
    self._headers = []

# ヘッダー1行毎に複数回呼ばれる。
def header(self, name, value):
    # ヘッダーの情報を追加する。
    self._headers.append([name, value])

# 本文を全て処理した時点で呼ばれる。
def end_of_message(self):
    header_indexes = {}

    # 各ヘッダーに順番に置換処理を行う。
    for name, value in self._headers:
        # ヘッダーのインデックス(何個目の同じ名前のヘッダーか)をカウントする。
        if name not in header_indexes:
            header_indexes[name] = 0
        header_indexes[name] += 1

        for pattern, replaced in self._patterns.items():
            replaced_value, _ = pattern.subn(replaced, value)
            # 置換した結果が元の値と異なるのであれば、ヘッダーを更新する。
            if value != replaced_value:
                self._change_header(name,
                                    header_indexes[name],
                                    replaced_value)
                break

# 1つのSMTPセッションの中で複数メールが送信される場合に、その切れ目で呼ばれる。
# 初期化(super().__init__())時にも呼ばれる。
# 1メール毎にセッションを分けるか、1セッションの中で複数メールを送るかは
# クライアント次第。
def reset(self):
    # 本文を管理するためのメンバー変数を初期化する。
    self._body = b""

# 本文。ある程度大きい場合は、チャンクに分かれて複数回呼ばれる。
def body(self, chunk):
    # 本文を追加する。
    self._body += chunk

# 本文を全て処理した時点で呼ばれる。
def end_of_message(self):
    for pattern, replaced in self._patterns.items():
        body = self._body.decode("utf-8")
        replaced_body, _ = pattern.subn(replaced, body)
        # 置換した結果が元の値と異なるのであれば、本文を更新する。
        if body != replaced_body:
            self._replace_body(replaced_body)

# 1つのSMTPセッションの中で複数メールが送信される場合に、その切れ目で呼ばれる。
# 初期化(super().__init__())時にも呼ばれる。
# 1メール毎にセッションを分けるか、1セッションの中で複数メールを送るかは
# クライアント次第。
def reset(self):
    self._body = b""

# ヘッダー1行毎に複数回呼ばれる。
def header(self, name, value):
    # Fromに"@clear-code.com"で終わるメールアドレスがあった場合に
    # 後続の処理をスキップしてメールを通過させる。
    if name.lower() == "From".lower():
        if value.endswith("@clear-code.com"):
            self._accept()

# 本文。ある程度大きい場合は、チャンクに分かれて複数回呼ばれる。
def body(self, chunk):
    self._body += chunk

# 本文を全て処理した時点で呼ばれる。
def end_of_message(self):
    for pattern, replaced in self._patterns.items():
        body = self._body.decode("utf-8")
        replaced_body, _ = pattern.subn(replaced, body)
        if body != replaced_body:
            self._replace_body(replaced_body)

# ヘッダー1行毎に複数回呼ばれる。
def header(self, name, value):
    # Fromに"@clear-code.com"で終わるメールアドレスがあった場合に
    # このメールを拒否する。
    if name.lower() == "From".lower():
        if value.endswith("@clear-code.com"):
            self._reject()

# 1つのSMTPセッションの中で複数メールが送信される場合に、その切れ目で呼ばれる。
# 初期化(super().__init__())時にも呼ばれる。
# 1メール毎にセッションを分けるか、1セッションの中で複数メールを送るかは
# クライアント次第。
def reset(self):
    self._headers = []

# ヘッダー1行毎に複数回呼ばれる。
def header(self, name, value):
    self._headers.append([name, value])

# 本文を全て処理した時点で呼ばれる。
def end_of_message(self):
    for name, value in self._headers:
        # Fromに"@clear-code.com"で終わるメールアドレスがあった場合に
        # このメールを隔離する。
        if name.lower() != "From".lower():
            continue
        if not value.endswith("@clear-code.com"):
            continue
        self._quarantine("This mail is quarantined " +
                         "since it has the header " +
                         f"{name}:{value}")
        break

#!/usr/bin/env python3

import milter.client

class MilterSetting(milter.client.Session):
    def __init__(self, context, config):
        super().__init__(context)
        milter_setting_section = config["milter_setting"]
        self._option = milter_setting_section["option_name"]


command_line = milter.client.CommandLine()
with command_line.run() as (client, options):
    client.register(MilterSetting, options.config)

#!/usr/bin/env python3

import milter.core
import milter.client

class MilterLogging(milter.client.Session):
    def __init__(self, context, config):
        super().__init__(context)
        self._logger = milter.core.Logger.default

    def body(self, chunk):
        self._logger.info(f"body: {chunk}")


command_line = milter.client.CommandLine()
with command_line.run() as (client, options):
    client.register(MilterLogging, options.config)

#!/usr/bin/env python3
#
# Copyright (C) 2022  Sutou Kouhei <kou@clear-code.com>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import subprocess
import sys

import milter.client

class MilterExternal(milter.client.Session):
    def __init__(self, context, timeout):
        super().__init__(context)
        self._timeout = timeout

    def end_of_message(self):
        command_line = [
            sys.executable,
            "-c",
            f"import time; time.sleep({self._timeout})",
        ]
        # 外部コマンドを別プロセスで実行する。
        process = subprocess.Popen(command_line)
        # 外部コマンドが完了した時点で呼ばれるコールバック関数を定義する。
        def on_exit(pid, wait_status):
            # コールバック関数を登録したイベントを削除する。
            self._remove_source(self._source_id)
            try:
                # 外部コマンド完了後の後処理部分。
                # ここでエラーが発生した場合も、 _context.emit() を呼ぶ必要がある。
                self._process_after_external_command()
            except Exception as exception:
                # _context.fallback_status を使うことで、
                # 設定したフォールバックステータスが効くようにする。
                self._context.emit("end_of_message_response",
                                   self._context.fallback_status)
        # コールバック関数を登録する。
        self._source_id = self._watch_child(process.pid, on_exit)
        # 外部コマンドの完了を待つ。
        # _context.emit() を実行して end_of_message が完了したことを
        # 知らせるまで待機する。
        self._delay_response()

    def _process_after_external_command(self):
        # 外部コマンド完了後の後処理を行う部分。
        print("Do process after external command")
        # 後処理が正常に完了したら、accept状態にして _context.emit() を呼ぶ。
        self._accept()
        self._context.emit("end_of_message_response",
                           self._context.status)

    def abort(self, status):
        if self._source_id is not None:
            self._remove_source(self._source_id)

    def reset(self):
        self._source_id = None

command_line = milter.client.CommandLine()
with command_line.run() as (client, options):
    client.register(MilterExternal, 3)