-- The Ada Milter API test milter
-- Copyright 2012 - 2013 B. Persson, Bjorn@Rombobeorn.se
--
-- This program is free software: you can redistribute it and/or modify it
-- under the terms of the GNU General Public License version 3, as published
-- by the Free Software Foundation.


with Ada.Strings.Fixed; use Ada.Strings.Fixed;
with Milter_API; use Milter_API;
with Berkeley_Exit_Codes;
with System_Log; use System_Log;
with Ada.Unchecked_Deallocation;
with Interfaces.C;
with Ada.Exceptions;
with Ada.Text_IO;

package body Test_Milter_Package is


   type Test_Action is (None, Test_Reject, Test_Discard, Test_Fail_Temporarily);


   type Message_Data is limited new Milter_Data with record
      Test_Message : Boolean;
      Action       : Test_Action;
   end record;
   type Message_Data_Pointer is access all Message_Data;
   procedure Free is new Ada.Unchecked_Deallocation(Message_Data,
                                                    Message_Data_Pointer);


   function Private_Data(Context : SMFICTX_Pointer) return Message_Data_Pointer
   is
   begin
      return Message_Data_Pointer(Milter_API.Private_Data(Context));
   end Private_Data;


   function Handle_Connection
      (Context        : SMFICTX_Pointer;
       Client_Name    : String;
       Client_Address : Sockaddr)
      return Action
   is
      Local_Client : Boolean := False;
      -- Allocate a message data record for this SMTP session.
      Data : constant Message_Data_Pointer := new Message_Data;
   begin
      Log(Debug, "Handle_Connection");
      if Milter_API.Private_Data(Context) /= null then
         Log(Warning,
             "The private data pointer isn't null in Handle_Connection. " &
             "Memory is probably leaking.");
      end if;
      -- Remember the pointer to the message data record.
      Set_Private_Data(Context, Milter_Data_Pointer(Data));
      declare
         use type Interfaces.Unsigned_8;
         Addr : constant IP_Address := Address(Client_Address);
      begin
         case Addr.Family is
            when IPv4 =>
               Local_Client := Addr.IPv4_Address(1) = 127;
            when IPv6 =>
               Local_Client := Addr.IPv6_Address = (1..15 => 0, 16 => 1);
         end case;
      exception
         when No_Address =>
            Log(Warning, "The MTA didn't provide the client's IP address.");
         when Unknown_Address_Type =>
            Log(Error, "The client address is of an unknown type.");
      end;
      Log(Debug,
          "client address: " & Address(Client_Address) &
          ", local client: " & Boolean'Image(Local_Client));
      if Local_Client then
         return Continue;
      else
         -- The test milter won't touch messages from this connection.
         return Accept_Definitely;
      end if;
   end Handle_Connection;


   function Handle_Helo
      (Context     : SMFICTX_Pointer;
       Stated_Name : String)
      return Action
   is
      Data : constant Message_Data_Pointer := Private_Data(Context);
   begin
      Log(Debug, "Handle_Helo");
      return Continue;
   end Handle_Helo;


   function Handle_Sender
      (Context   : SMFICTX_Pointer;
       Sender    : String;
       Arguments : Arguments_Handle)
      return Action
   is
      Data : constant Message_Data_Pointer := Private_Data(Context);
   begin
      Log(Debug, "Handle_Sender");
      -- Initialize the message data record, or clear it of data from the
      -- previous message in the SMTP session.
      Data.Test_Message := False;
      Data.Action := None;
      return Continue;
   end Handle_Sender;


   function Handle_Recipient
      (Context   : SMFICTX_Pointer;
       Recipient : String;
       Arguments : Arguments_Handle)
      return Action
   is
      Data : constant Message_Data_Pointer := Private_Data(Context);
   begin
      Log(Debug, "Handle_Recipient " & Recipient);
      if Index(Recipient, "Ada_Milter_API_test_milter") /= 0 then
         Data.Test_Message := True;
      end if;
      return Continue;
   end Handle_Recipient;


   function Handle_Data(Context : SMFICTX_Pointer) return Action is
      Data : constant Message_Data_Pointer := Private_Data(Context);
   begin
      Log(Debug, "Handle_Data");
      if Data.Test_Message then
         return Continue;
      else
         -- This message is not intended for the test milter.
         return Accept_Definitely;
      end if;
   end Handle_Data;


   function Handle_Unknown_Command
      (Context : SMFICTX_Pointer;
       Command : String)
      return Action
   is
      Data : constant Message_Data_Pointer := Private_Data(Context);
   begin
      Log(Debug, "Handle_Unknown_Command");
      return Continue;
   end Handle_Unknown_Command;


   function Handle_Header
      (Context : SMFICTX_Pointer;
       Name    : String;
       Value   : String)
      return Action
   is
      Data : constant Message_Data_Pointer := Private_Data(Context);
   begin
      Log(Debug, "Handle_Header " & Name);
      return Continue;
   end Handle_Header;


   function Handle_End_Of_Headers(Context : SMFICTX_Pointer) return Action is
      Data : constant Message_Data_Pointer := Private_Data(Context);
   begin
      Log(Debug, "Handle_End_Of_Headers");
      return Continue;
   end Handle_End_Of_Headers;


   function Handle_Body
      (Context    : SMFICTX_Pointer;
       Body_Chunk : String)
      return Action
   is
      Data : constant Message_Data_Pointer := Private_Data(Context);
   begin
      Log(Debug, "Handle_Body");
      return Continue;
   end Handle_Body;


   function Handle_End_Of_Message(Context : SMFICTX_Pointer) return Action is
      Data : constant Message_Data_Pointer := Private_Data(Context);
   begin
      Log(Debug, "Handle_End_Of_Message");
      return Reject;
   end Handle_End_Of_Message;


   procedure Handle_Abort(Context : SMFICTX_Pointer) is
      Data : constant Message_Data_Pointer := Private_Data(Context);
   begin
      Log(Debug, "Handle_Abort");
   end Handle_Abort;


   procedure Handle_Close(Context : SMFICTX_Pointer) is
      Data : Message_Data_Pointer := Private_Data(Context);
   begin
      Log(Debug, "Handle_Close");
      -- Deallocate the message data record.
      Free(Data);
      Set_Private_Data(Context, null);
   end Handle_Close;


   function Run return Ada.Command_Line.Exit_Status is
      use Ada.Exceptions;
      use Berkeley_Exit_Codes;
      use Ada.Text_IO;

      function umask(mask : Interfaces.C.unsigned) return Interfaces.C.unsigned;
      pragma Import(C, umask);
      Dummy : Interfaces.C.unsigned;  -- to soak up the result from umask
      pragma Unreferenced(Dummy);

      Socket_Name : constant String := "/var/spool/test_milter/milter_socket";
   begin
      Log(Info,
          "Starting, using " &
          "Ada Milter API " & Milter_API.Binding_Version_String & " and " &
          "Libmilter " & Milter_API.Libmilter_Version_String & '.');
      Register(Name                         => "test_milter/Libmilter",
               Connected                    => Handle_Connection'Access,
               Helo                         => Handle_Helo'Access,
               Sender                       => Handle_Sender'Access,
               Recipient                    => Handle_Recipient'Access,
               Data                         => Handle_Data'Access,
               Unknown_Command              => Handle_Unknown_Command'Access,
               Header                       => Handle_Header'Access,
               End_Of_Headers               => Handle_End_Of_Headers'Access,
               Body_Chunk                   => Handle_Body'Access,
               End_Of_Message               => Handle_End_Of_Message'Access,
               Aborted                      => Handle_Abort'Access,
               Closed                       => Handle_Close'Access);
      -- Clear the permissions mask to allow the MTA to use the socket.
      Dummy := umask(0);
      -- Tell Libmilter where to create the socket.
      Set_Socket("unix:" & Socket_Name);
      -- Open the socket. Delete any existing socket, assuming it was left
      -- behind because of a crash.
      begin
         Open_Socket(Remove_Old_Socket => True);
      exception
         when Milter_API.Failure =>
            -- Libmilter has already logged sufficient error messages.
            return Cannot_Create_File;
      end;
      -- Hand over control to Libmilter.
      Milter_API.Main;
      return Ada.Command_Line.Success;
   exception
      when E : Milter_API.Failure =>
         Log(Error, Exception_Message(E));
         return Ada.Command_Line.Failure;
      when E : Milter_API.Unknown_Error =>
         Log(Error, Exception_Message(E));
         return Software_Error;
      when E : others =>
         Put_Line(Standard_Error, Exception_Information(E));
         Log(Error,
             "Unexpected error: " & Exception_Name(E) & ": " &
             Exception_Message(E));
         return Software_Error;
   end Run;


end Test_Milter_Package;