Source file received.ml

open Colombe
open Mrmime

module Option = struct
  let map f = function Some x -> Some (f x) | None -> None
end

type 'a stream = unit -> 'a option

let error_msgf fmt = Format.kasprintf (fun msg -> Error (`Msg msg)) fmt

type protocol = [ `ESMTP | `SMTP | `Atom of string ]

(* TODO(dinosaure): according to RFC3848, we should add:
   - ESMTPA
   - ESMTPS
   - ESMTPSA
   - LMTP
   - LMTPA
   - LMTPSA
*)
type link = [ `TCP | `Atom of string ]

type 'a with_info = Only of 'a | With of 'a * info
and info = [ `Address of Domain.t | `Domain_and_address of Domain.t * Domain.t ]

type t = {
  from : Domain.t with_info option;
  by : Domain.t with_info option;
  via : link option;
  _with : protocol option;
  id :
    [ `Local of Emile.local | `MsgID of Mrmime.MessageID.t | `Atom of string ]
    option;
  _for : Path.t option;
  date_time : Mrmime.Date.t;
}

let received_with { _with; _ } = _with
let received_via { via; _ } = via
let id { id; _ } = id
let date_time { date_time; _ } = date_time

let map_domain = function
  | Domain.IPv4 v -> `Addr (Emile.IPv4 v)
  | Domain.IPv6 v -> `Addr (Emile.IPv6 v)
  | Domain.Extension (ldh, v) -> `Addr (Emile.Ext (ldh, v))
  | Domain.Domain v -> `Domain v

let msg_id t =
  let domain =
    match t.by with
    | Some (Only domain) -> Some domain
    | Some (With (domain, _)) -> Some domain
    | None -> None in
  let domain = Option.map map_domain domain in
  match (t.id, domain) with
  | Some (`Local local), Some domain -> Some (local, domain)
  | Some (`Atom atom), Some domain -> Some ([ `Atom atom ], domain)
  | Some (`MsgID v), _ -> Some (v :> Emile.local * Emile.domain)
  | _ -> None

let equal t0 t1 =
  let msg_id0 = msg_id t0 in
  let msg_id1 = msg_id t1 in
  match (msg_id0, msg_id1) with
  | Some (local0, domain0), Some (local1, domain1) ->
      Emile.equal_local local0 local1 && Emile.equal_domain domain0 domain1
  | Some _, None | None, Some _ | None, None -> false

let compare t0 t1 =
  let sup = 1 and inf = -1 in
  let msg_id0 = msg_id t0 in
  let msg_id1 = msg_id t1 in
  match (msg_id0, msg_id1) with
  | Some (local0, domain0), Some (local1, domain1) ->
      let res = Emile.compare_local local0 local1 in
      if res = 0 then Emile.compare_domain domain0 domain1 else res
  | Some _, None -> sup
  | None, Some _ -> inf
  | None, None -> 0

let received_by { by; _ } = by
let received_from { from; _ } = from
let received_for { _for; _ } = _for
let tcp = `TCP
let link v = `Atom v
let smtp = `SMTP
let esmtp = `ESMTP
let protocol v = `Atom v

let make ?from ?by ?via ?protocol ?id _for ~zone ptime =
  let date_time = Date.of_ptime ~zone ptime in
  let id = match id with Some id -> Some (`MsgID id) | None -> None in
  { from; by; via; id; _with = protocol; _for; date_time }

let ( <.> ) f g x = f (g x)

let pp_extended_domain ppf = function
  | Only v -> Domain.pp ppf v
  | With (v, `Address addr) -> Fmt.pf ppf "%a (%a)" Domain.pp v Domain.pp addr
  | With (v, `Domain_and_address (domain, addr)) ->
      Fmt.pf ppf "%a (%a %a)" Domain.pp v Domain.pp domain Domain.pp addr

let pp_via ppf = function
  | `Atom v -> Fmt.string ppf v
  | `TCP -> Fmt.string ppf "tcp"

let pp_with ppf = function
  | `ESMTP -> Fmt.string ppf "esmtp"
  | `SMTP -> Fmt.string ppf "smtp"
  | `Atom v -> Fmt.string ppf v

let pp_id ppf = function
  | `MsgID v -> MessageID.pp ppf v
  | `Local v -> Emile.pp_local ppf v
  | `Atom v -> Fmt.string ppf v

let pp_for = Path.pp

let pp ppf t =
  let ptime, tz_offset_s = (Result.get_ok <.> Date.to_ptime) t.date_time in
  Fmt.pf ppf
    "from: %a@\n\
     by:   %a@\n\
     via:  %a@\n\
     with: %a@\n\
     id:   %a@\n\
     for:  %a@\n\
     date: %a@\n"
    Fmt.(option pp_extended_domain)
    t.from
    Fmt.(option pp_extended_domain)
    t.by
    Fmt.(option pp_via)
    t.via
    Fmt.(option pp_with)
    t._with
    Fmt.(option pp_id)
    t.id
    Fmt.(option pp_for)
    t._for
    (Ptime.pp_rfc3339 ~tz_offset_s ())
    ptime

let some x : _ option = Some x

module Decoder = struct
  open Angstrom

  let is_wsp = function ' ' | '\t' -> true | _ -> false

  let domain =
    Domain.Decoder.ipv4_address_literal
    >>| (fun v -> Domain.IPv4 v)
    <|> (Domain.Decoder.ipv6_addr >>| fun v -> Domain.IPv6 v)
    <|> Domain.Decoder.domain

  let address_literal = Domain.Decoder.address_literal
  let atom = lift (fun x -> `Atom x) Path.Decoder.atom
  let date_time = Date.Decoder.date_time
  let fws = skip_while is_wsp
  let cfws = Emile.Parser.cfws

  let tcp_info =
    address_literal
    >>= (fun v -> return (`Address v))
    <|> ( domain >>= fun domain ->
          skip_while is_wsp *> address_literal >>= fun v ->
          return (`Domain_and_address (domain, v)) )

  let via = string_ci "via"
  let _with = string_ci "with"
  let id = string_ci "id"
  let _for = string_ci "for"
  let link = atom >>= function `Atom "TCP" -> return `TCP | v -> return v

  let protocol =
    atom >>= function
    | `Atom "ESMTP" -> return `ESMTP
    | `Atom "SMTP" -> return `SMTP
    | v -> return v

  let via = option () cfws *> via *> fws *> link
  let via = lift some via
  let _with = option () cfws *> _with *> fws *> protocol
  let _with = lift some _with

  let id =
    let msg_id = MessageID.Decoder.message_id >>| fun v -> `MsgID v in
    let local =
      Path.Decoder.local_part >>| function
      | `String v -> `Local [ `String v ]
      | `Dot_string vs -> List.map (fun v -> `Atom v) vs |> fun vs -> `Local vs
    in
    option () cfws *> id *> fws *> (msg_id <|> local <|> atom)

  let id = lift some id

  let _for =
    option () cfws
    *> _for
    *> fws
    *> (Path.Decoder.path
       >>= (fun v -> return v)
       <|> ( Path.Decoder.mailbox >>= fun (local, domain) ->
             return { Path.local; domain; rest = [] } ))

  let _for = lift some _for

  let extended_domain =
    domain
    >>= (fun domain ->
          skip_while is_wsp *> char '(' *> tcp_info <* char ')'
          >>= fun tcp_info -> return (With (domain, tcp_info)))
    <|> ( address_literal >>= fun domain ->
          skip_while is_wsp *> char '(' *> tcp_info <* char ')'
          >>= fun tcp_info -> return (With (domain, tcp_info)) )
    <|> (domain >>= fun domain -> return (Only domain))

  let from_domain = string_ci "from" *> skip_while is_wsp *> extended_domain

  let by_domain =
    option () cfws *> string_ci "by" *> skip_while is_wsp *> extended_domain

  let separator =
    option () cfws *> char ';' *> skip_while is_wsp
    <|> (take_while1 is_wsp >>= fun _ -> return ())

  let to_end_of_input buf =
    fix @@ fun m ->
    at_end_of_input >>= function
    | false ->
        available >>= take >>= fun str ->
        Buffer.add_string buf str ;
        m
    | true -> return (Buffer.contents buf)

  let date_time =
    date_time
    <|> ( to_end_of_input (Buffer.create 0x100) >>= fun str ->
          match Ptime.of_rfc3339 ~sub:true (String.trim str) with
          | Ok (ptime, Some tz_offset_s, _) ->
              let hh = tz_offset_s / 3600 in
              let mm = tz_offset_s mod 3600 / 60 in
              let zone = Date.Zone.TZ (hh, mm) in
              let date_time = Date.of_ptime ~zone ptime in
              return date_time
          | Ok (ptime, None, _) ->
              let zone = Date.Zone.UT (* TODO *) in
              let date_time = Date.of_ptime ~zone ptime in
              return date_time
          | _ -> fail "date_time" )

  let stamp =
    option None (lift some (fws *> from_domain)) >>= fun from_domain ->
    option None (lift some (fws *> by_domain)) >>= fun by_domain ->
    option None (fws *> via) >>= fun via ->
    option None (fws *> _with) >>= fun _with ->
    option None (fws *> id) >>= fun id ->
    option None (fws *> _for) >>= fun _for ->
    separator >>= fun _ ->
    date_time >>= fun date_time ->
    return
      { from = from_domain; by = by_domain; via; _with; id; _for; date_time }
end

module Encoder = struct
  open Prettym

  let extended_domain ppf = function
    | Only domain -> string ppf (Domain.to_string domain)
    | With (domain, `Address addr) ->
        eval ppf
          [ !!string; spaces 1; char $ '('; !!string; char $ ')' ]
          (Domain.to_string domain) (Domain.to_string addr)
    | With (domain, `Domain_and_address (v, addr)) ->
        eval ppf
          [
            !!string;
            spaces 1;
            char $ '(';
            !!string;
            spaces 1;
            !!string;
            char $ ')';
          ]
          (Domain.to_string domain) (Domain.to_string v) (Domain.to_string addr)

  let via ppf = function `TCP -> string ppf "tcp" | `Atom v -> string ppf v

  let via ppf = function
    | None -> ppf
    | Some v -> eval ppf [ string $ "via"; spaces 1; !!via ] v

  let _with ppf = function
    | `ESMTP -> string ppf "esmtp"
    | `SMTP -> string ppf "smtp"
    | `Atom v -> string ppf v

  let _with ppf = function
    | None -> ppf
    | Some v -> eval ppf [ string $ "with"; spaces 1; !!_with ] v

  let id ppf = function
    | `MsgID v -> Mrmime.MessageID.Encoder.message_id ppf v
    | `Local v -> string ppf (Fmt.str "%a" Emile.pp_local v)
    | `Atom v -> string ppf v

  let id ppf = function
    | None -> ppf
    | Some v -> eval ppf [ box; string $ "id"; spaces 1; !!id; close ] v

  let _for ppf v = string ppf (Path.Encoder.to_string v)

  let _for ppf = function
    | None -> ppf
    | Some v -> eval ppf [ box; string $ "for"; spaces 1; !!_for; close ] v

  let opt_info ppf (v_via, v_with, v_id, v_for) =
    eval ppf
      [
        box;
        !!via;
        fws;
        !!_with;
        fws;
        !!id;
        fws;
        !!_for;
        char $ ';';
        close;
        new_line;
      ]
      v_via v_with v_id v_for

  let from_domain ppf = function
    | Some v ->
        eval ppf
          [ box; string $ "from"; fws; !!extended_domain; close; new_line ]
          v
    | None -> ppf

  let by_domain ppf = function
    | Some v ->
        eval ppf
          [ box; string $ "by"; fws; !!extended_domain; close; new_line ]
          v
    | None -> ppf

  let received ppf { from; by; via; _with; id; _for; date_time } =
    eval ppf
      [
        box;
        !!from_domain;
        !!by_domain;
        !!opt_info;
        !!Mrmime.Date.Encoder.date;
        close;
      ]
      from by (via, _with, id, _for) date_time

  let as_field ppf v =
    eval ppf
      [
        string $ "Received";
        char $ ':';
        tbox 1;
        spaces 1;
        !!received;
        close;
        new_line;
      ]
      v
end

let received = Field_name.v "Received"

let of_unstructured v =
  let v =
    List.fold_left
      (fun a -> function #Unstrctrd.elt as x -> x :: a | _ -> a)
      [] v in
  let v = List.rev v in
  match Unstrctrd.of_list v with
  | Ok v ->
      let v = Unstrctrd.fold_fws v in
      Unstrctrd.to_utf_8_string v
  | Error _ -> assert false
(* TODO *)

let of_stream stream =
  let decoder = Hd.decoder Field_name.Map.empty in
  let rec go acc =
    match Hd.decode decoder with
    | `Field field -> (
        let (Field.Field (field_name, w, v)) = Location.prj field in
        match w with
        | Field.Unstructured ->
            if Field_name.equal field_name received
            then
              let v = of_unstructured v in
              match Angstrom.parse_string ~consume:Prefix Decoder.stamp v with
              | Ok v -> go (v :: acc)
              | Error _err -> go acc
            else go acc
        | _ -> go acc)
    | `Malformed err -> error_msgf "%s" err
    | `End rest -> Ok (rest, List.rev acc)
    | `Await ->
        let buf, off, len =
          match stream () with
          | Some (buf, off, len) -> (buf, off, len)
          | None -> ("", 0, 0) in
        Hd.src decoder buf off len ;
        go acc in
  go []