open Unix;
open Pcre;
open Hashtbl;

exception MatchedOk;
exception MatchedSpam;

type state =
  [ Helo
  | Mail
  | Rcpt
  | Data
  | Datacmd
  | Datadone
  | Quit
  | Unknown ]
;
type svstate =
  [ PreliminaryOk
  | Ok
  | IntermediateOk
  | TransientError
  | PermanentError
  | ServUnknown ]
;

value procmailrc = ".procmailrc-filter";
value spawn = "spawn";
value filterargs = ["/openpkg/bin/procmail"; procmailrc];

value usrre = regexp ": *<?([^<>@]+)(@[\\w\.-])?";

value nonspam = ["X-Whitelist: match"; "X-DSPAM-Result: Innocent"];
value spam = ["X-DSPAM-Result: Spam"];

value postmap = ["postmap"; "-q"];
value aliastype = "hash";
value aliasfile = "/openpkg/etc/postfix/aliases";
value aliases = [aliastype ^ ":" ^ aliasfile];

value svstate_str s =
  match s with
  [ PreliminaryOk -> "PreliminaryOk"
  | Ok -> "Ok"
  | IntermediateOk -> "IntermediateOk"
  | TransientError -> "TransientError"
  | PermanentError -> "PermanentError"
  | ServUnknown -> "Unknown" ]
;

value clstate_str s =
  match s with
  [ Helo -> "helo"
  | Mail -> "mail"
  | Rcpt -> "rcpt"
  | Datacmd -> "datacmd"
  | Data -> "data"
  | Datadone -> "datadone"
  | Quit -> "quit"
  | Unknown -> "unknown" ]
;

value hashkeys h =
  let getkey a b = \@ [a] in
  Hashtbl.fold getkey h []
;

value rec resolvealias user =
  let inc =
    Unix.open_process_in (String.concat " " (postmap @ [user] @ aliases))
  in
  let out = try resolvealias (input_line inc) with [ End_of_file -> user ] in
  do { ignore (close_process_in inc); out }
;

value removecr s =
  let len = String.length s in
  try if s.[len - 1] = '\r' then String.sub s 0 (len - 1) else s with _ -> s
; 

value get_addr s = (Unix.gethostbyname s).Unix.h_addr_list.(0);

value main_server serv_fun =
  if Array.length Sys.argv < 5 then
    Printf.eprintf "usage: %s host port host port\n" Sys.argv.(0)
  else
    try
      let port = int_of_string Sys.argv.(2) in
      let my_address = get_addr Sys.argv.(1) in
      establish_server serv_fun (Unix.ADDR_INET my_address port)
    with
    [ Failure "int_of_string" -> Printf.eprintf "bad port number\n" ]
;

value proxy_service clientin clientout =
  try
    let my_address = get_addr Sys.argv.(3) in
    let port = int_of_string Sys.argv.(4) in
    let (svin, svout) = open_connection (Unix.ADDR_INET my_address port) in
    let clstate = ref Unknown in
    let svstate = ref ServUnknown in
    let mail = ref "" in
    let userhash = Hashtbl.create 10 in
    let lines = ref [] in
    let first = ref True in
    let isspam = ref False in

    let tokcmp s1 s2 =
      try s1 = String.lowercase (String.sub s2 0 (String.length s1)) with
      [ Invalid_argument "String.sub" -> False ]
    in

    let update_clstate s =
      do {
        if clstate.val = Datacmd then clstate.val := Data else ();
        if tokcmp "quit" s then clstate.val := Quit else ();
        if tokcmp "helo " s then clstate.val := Helo else ();
        if tokcmp "ehlo " s then clstate.val := Helo else ();
        if tokcmp "mail from:" s then clstate.val := Mail else ();
        if clstate.val = Mail && tokcmp "rcpt to:" s then clstate.val := Rcpt
        else ();
        if tokcmp "rset" s then do {
          Hashtbl.clear userhash;
          clstate.val := Unknown;
          lines.val := [];
          mail.val := "";
          first.val := True;
          isspam.val := False
        }
        else ();
        if clstate.val = Rcpt && tokcmp "data" s then clstate.val := Datacmd
        else ();
        if clstate.val = Data && s = "." then clstate.val := Datadone else ()
      }
    in

    let dbglog fmt =
      let a x =
        Printf.printf (\^^ (format_of_string "%d ") (format_of_string fmt))
          (Unix.getpid ()) x
      in
      do { flush Pervasives.stdout; a }
    in

    let cllog str = dbglog "cl: %s %s\n" (clstate_str clstate.val) str in
    let svlog str = dbglog "sv: %s %s\n" (svstate_str svstate.val) str in
    let prlog fmt x =
      dbglog (\^^ (format_of_string "pr: ") (format_of_string fmt)) x
    in

    let update_svstate s =
      do {
        if String.length s >= 4 then do {
          if s.[0] = '1' then svstate.val := PreliminaryOk else ();
          if s.[0] = '2' then svstate.val := Ok else ();
          if s.[0] = '3' then svstate.val := IntermediateOk else ();
          if s.[0] = '4' then svstate.val := TransientError else ();
          if s.[0] = '5' then svstate.val := PermanentError else ()
        }
        else ();
        svlog s;
        s
      }
    in

    let svread () =
      do {
        clear_nonblock (descr_of_in_channel svin);
        let a = Buffer.create 1000 in
        try
          while True do {
            Buffer.add_string a
              (update_svstate (removecr (input_line svin)) ^ "\r\n");
            set_nonblock (descr_of_in_channel svin)
          }
        with
        [ Sys_blocked_io -> ()
        | End_of_file -> () ];
        Buffer.contents a
      }
    in

    do {
      output_string clientout (svread ());
      flush clientout;
      while clstate.val <> Quit do {
        let cl = removecr (input_line clientin) in
        update_clstate cl;
        cllog cl;
        match clstate.val with
        [ Mail ->
            do {
              mail.val := cl;
              prlog "%s\n" "faking mail response";
              output_string clientout "250 Ok\r\n";
              flush clientout
            }
        | Rcpt ->
            try
              let getuser str = (extract ~rex:usrre str).(1) in
              let user = resolvealias (String.lowercase (getuser cl)) in
              do {
                prlog "user: %s\n" user;
                Hashtbl.replace userhash user cl;
                prlog "userlist: %s\n"
                  (String.concat ", " (hashkeys userhash));
                prlog "%s\n" "faking rcpt response to client";
                output_string clientout "250 Ok\r\n";
                flush clientout
              }
            with
            [ Not_found -> prlog "%s\n" "no user" ]
        | Datacmd ->
            do {
              prlog "%s\n" "faking data response to client";
              output_string clientout
                "354 End data with <CR><LF>.<CR><LF>\r\n";
              flush clientout
            }
        | Data -> lines.val := lines.val @ [cl]
        | Datadone ->
            do {
              clstate.val := Unknown;
              let datatoserver lines =
                do {
                  output_string svout
                    (String.concat "\r\n" (lines @ [".\r\n"]));
                  flush svout;
                  if first.val then do {
                    first.val := False;

                    if isspam.val then do {
                      ignore (svread ());
                      prlog "%s\n" "isspam: faking failed to client";
                      output_string clientout "550 Failed\r\n";
                    } else do {
                      prlog "%s\n" "isspam: not spam";
                      output_string clientout (svread ());
                    };
                    flush clientout
                  }
                  else ignore (svread ())
                }
              in
              let fakestuff user =
                do {
                  prlog "sending fake mail cmd: %s\n" mail.val;
                  output_string svout (mail.val ^ "\r\n");
                  flush svout;
                  ignore (svread ());
                  prlog "sending fake rcpt cmd: %s\n"
                    (Hashtbl.find userhash user);
                  output_string svout (Hashtbl.find userhash user ^ "\r\n");
                  flush svout;
                  ignore (svread ());
                  prlog "%s\n" "sending fake data cmd";
                  output_string svout "DATA\r\n";
                  flush svout;
                  ignore (svread ())
                }
              in
              let filtermail user =
                try
                  let pwent = getpwnam user in
                  do {
                    close_in (open_in (pwent.pw_dir ^ "/" ^ procmailrc));
                    prlog "%s\n" "message read, spawning filter";
                    let (inc, outc) =
                      Unix.open_process
                        (String.concat " " [spawn; user :: filterargs])
                    in
                    output_string outc (String.concat "\n" lines.val);
                    close_out outc;
                    fakestuff user;

                    let rec flines inc =
                      let line = input_line inc in
                      let checkline ex a = 
                        if line = a then raise ex
                        else ()
                      in

                      try do {
                        List.iter (checkline MatchedOk) nonspam;
                        List.iter (checkline MatchedSpam) spam;
                        List.iter (checkline MatchedOk) [""];
                        [ line::flines inc ]
                      } with
                      [ MatchedOk -> [ line ] 
                      | MatchedSpam -> do { isspam.val := True; [ line ] } ]
                    in

                    let svlog ln = do {
(*
                      prlog "flines: %s\n" ln;
*)
                      output_string svout (ln^"\r\n");
                    } in

                    List.iter svlog (flines inc); 

                    try
                      while True do {
                        output_string svout ((input_line inc)^"\r\n");
                      }
                    with [ End_of_file -> () ];

                    datatoserver [];
  
                    ignore
                      (match close_process (inc, outc) with
                       [ WEXITED x -> prlog "exited %d\n" x
                       | WSIGNALED x -> prlog "signaled %d\n" x
                       | _ -> prlog "%s\n" "unknown" ]);

                    prlog "%s\n" "filter done, sending RSET";
                    clstate.val := Unknown;
                    output_string svout "RSET\n";
                    flush svout;
                    prlog "%s\n" "sent RSET";
                    ignore (svread ())
                  }
                with
                [ Not_found ->
                    do {
                      prlog "%s\n" "user not found";
                      fakestuff user;
                      datatoserver lines.val
                    }
                | Sys_error x ->
                    do {
                      prlog "procmailrc not found: %s\n" x;
                      fakestuff user;
                      datatoserver lines.val
                    }
                | _ ->
                    do {
                      prlog "%s\n" "filter failed";
                      fakestuff user;
                      datatoserver lines.val
                    } ]
              in
              List.iter filtermail (hashkeys userhash)
            }
        | _ ->
            do {
              output_string svout (cl ^ "\r\n");
              flush svout;
              output_string clientout (svread ());
              flush clientout
            } ]
      }
    }
  with
  [ Sys_error x -> do { Printf.printf "%s\n" x; exit 0 }
  | End_of_file -> do { Printf.printf "End_of_file\n"; exit 0 } ]
;

ignore (Unix.setsid ());
Unix.handle_unix_error main_server proxy_service;