#include <arpa/inet.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>

#include "buffer.h"
#include "case.h"
#include "constmap.h"
#include "error.h"
#include "exit.h"
#include "fd.h"
#include "fmt.h"
#include "logmsg.h"
#include "open.h"
#include "pathexec.h"
#include "sig.h"
#include "str.h"
#include "stralloc.h"
#include "uint_t.h"
#include "wait.h"

#include "control.h"
#include "fmtqfn.h"
#include "qmail.h"
#include "rcpthosts.h"

#ifdef USE_CONFIG
  #include "fehsqm-config.h"
#else
  #include "auto_qmail.h"
#endif

#define WHO "qmail-dksign"

#define DOMAINKEYS "ssl/domainkeys/"

/**
  @file qmail-dksign.c -- generate signature and attach in DKIM header to outgoing message

  Steps:
  ------
  a) DKIM controls:  get private key for sending domain
  b) Prepare two staging files at queue/dkim (before and after signing)
  c) Read input at fd0 and insert CR for every line and store at dkim/x/pre
  d) DKIM sign the message with provided private key and store at dkim/y/post
  e) Copy signed file from fd to 0
  f) Invoke qmail-remote (respecting the \r\n)
  g) Remove staging files (pre/post)

  Hack for hybrid signatures:
  ---------------------------

  a) selector is a link to RSA private key
  b) selector2 is a link to Ed25519 private key
  c) Both are provided in the 'selector' field of dkimdomains separated by colon
  d) The coupled selector information is provided to qmail-dkim as: -yselector ,-Yselector2
  e) The RSA privat key is given unaltered
  f) The Ed25519 private is supplied as additional argument
*/

char bufin[1000];  // RFC 5322: 998 chars - why?
buffer bi = BUFFER_INIT(read, 0, bufin, sizeof(bufin));
char bufout[1000];
buffer bo = BUFFER_INIT(write, 1, bufout, sizeof(bufout));

void die(int e)
{
  _exit(e);
}

void die_write(char *fn)
{
  unlink(fn);
  die(53);
}

void die_read()
{
  die(54);
}

void out(char *s)
{
  if (buffer_puts(&bo, s) == -1) _exit(111);
}

void zero()
{
  if (buffer_put(&bo, "\0", 1) == -1) _exit(111);
}

void zerodie()
{
  zero();
  buffer_flush(&bo);
  _exit(111);
}

stralloc fndkin = {0};
stralloc fndkout = {0};

stralloc sender = {0};  // will be re-written
stralloc senddomain = {0};
stralloc originator = {0};
stralloc dkimdomains = {0};
struct constmap mapdkimdomains;

stralloc ecckey = {0};
stralloc rsakey = {0};
char *dkimparams = 0;

void temp_nomem()
{
  out("ZOut of memory. (#4.3.0)\n");
  zerodie();
}

void temp_chdir()
{
  out("ZUnable to switch to target directory. (#4.3.0)\n");
  zerodie();
}

void temp_create()
{
  out("ZUnable to create DKIM stage file: ");
  out(error_str(errno));
  out(fndkin.s);
  out(". (#4.3.0)\n");
  zerodie();
}

void temp_unlink()
{
  out("ZUnable to unlink DKIM stage file. (#4.3.0)\n");
  zerodie();
}

void temp_control()
{
  out("ZUnable to read DKIM control files. (#4.3.0)\n");
  zerodie();
}

void perm_usage()
{
  out("Zqmail-dksign was invoked improperly. (#5.3.5)\n");
  zerodie();
}

void temp_read()
{
  out("DUnable to read message for DKIM signing. (#4.3.0)\n");
  zerodie();
}

void temp_nosignkey()
{
  out("DCan't read sign key: ");
  out(rsakey.s);
  out(" or ");
  out(ecckey.s);
  out(". (#4.3.0)\n");
  zerodie();
}

int get_controls()
{
  int i;
  stralloc domname = {0};

  if (control_init() == -1) temp_control();

  switch (control_readfile(&dkimdomains, "control/dkimdomains", 0)) {
    case -1: return 0;
    case 0:
      if (!constmap_init(&mapdkimdomains, "", 0, 1)) temp_nomem();
      break;
    case 1:
      if (!constmap_init(&mapdkimdomains, dkimdomains.s, dkimdomains.len, 1)) temp_nomem();
      break;
  }

  /* Check for disabled DKIM send domains */

  if (!stralloc_copys(&domname, "!")) temp_nomem();
  if (!stralloc_cats(&domname, senddomain.s)) temp_nomem();
  if (constmap(&mapdkimdomains, domname.s, domname.len)) return 0;

  /* Parenting domains; senddomain 0-terminated; lowercase */

  for (i = 0; i <= senddomain.len; ++i) {
    if ((i == 0) || (senddomain.s[i] == '.')) {
      if ((dkimparams = constmap(&mapdkimdomains, senddomain.s + i, senddomain.len - i - 1))) {
        if (!stralloc_copys(&sender, senddomain.s + i)) temp_nomem();
        if (!stralloc_0(&sender)) temp_nomem();
        return 3;
      }
    }
  }

  /* We sign only senddomains we take responsibility for: rcpthosts */

  if ((dkimparams = constmap(&mapdkimdomains, "=", 1))) {
    if (rcpthosts_init() == -1) temp_control();
    if (rcpthosts(originator.s, originator.len)) {
      if ((control_readline(&sender, "control/defaultdomain") != 1))
        if (control_readline(&sender, "control/me") == -1) temp_control();
      if (!stralloc_0(&sender)) temp_nomem();
      return 2;
    }
  }

  /* Default settings for MTA: 'defaultdomain' or even 'me' */

  if ((dkimparams = constmap(&mapdkimdomains, "*", 1))) {
    if ((control_readline(&sender, "control/defaultdomain") != 1))
      if (control_readline(&sender, "control/me") == -1) temp_control();
    if (!stralloc_0(&sender)) temp_nomem();
    return 1;
  }

  return 0;
}

void fnmake_dkim(unsigned long id)
{
  fndkin.len = fmtqfn(fndkin.s, "queue/dkim/", id, 1);
  id += id;
  fndkout.len = fmtqfn(fndkout.s, "queue/dkim/", id, 1);
}

void dkim_unlink()
{
  if (unlink(fndkin.s) == -1)
    if (errno != ENOENT) temp_unlink();
  if (unlink(fndkout.s) == -1)
    if (errno != ENOENT) temp_unlink();
}

void dkim_stage()
{
  int r;
  int fd;
  char ch;
  struct stat st;

  if (!stralloc_ready(&fndkin, FMTQFN)) temp_nomem();
  if (!stralloc_ready(&fndkout, FMTQFN)) temp_nomem();

  fnmake_dkim(getpid());  // pre-staging
  dkim_unlink();          // duplicate, left over file
  fd = open_excl(fndkin.s);
  if (fd == -1) die_write(fndkin.s);

  buffer_init(&bi, read, 0, bufin, sizeof(bufin));
  buffer_init(&bo, write, fd, bufout, sizeof(bufout));

  for (;;) {
    r = buffer_get(&bi, &ch, 1);
    if (r == 0) break;
    if (r == -1) temp_read();
    if (ch == '\r') continue;

    while (ch != '\n') {
      buffer_put(&bo, &ch, 1);
      r = buffer_get(&bi, &ch, 1);
      if (r == -1) temp_read();
    }
    buffer_put(&bo, "\r\n", 2);
  }

  if (buffer_flush(&bo) == -1) die(51);
  if (fstat(fd, &st) == -1) die_read();
  if (fsync(fd) == -1) die_write(fndkin.s);
  if (close(fd) == -1) die_write(fndkin.s);
}

/* to construct DKIM information */

stralloc selector = {0};
stralloc selectore = {0};
stralloc sdid = {0};
stralloc auid = {0};
stralloc expire = {0};
stralloc canon = {0};   // -c r = relax, s = simple, t = relaxed/simple, u = simple/realxed
stralloc hash = {0};    // -z 1/2/3/4/5 sha1/sha2/both/ed25519/ed25519+rsa-sha256
stralloc length = {0};  // -l

/**

  qmail-dkim [-h|-v|-s] [tags] <msgfile> [<RSAkeyfile> <outfile> <Ed25519keyfile>]
  --------------------------------------------------------------------------------
  tags:
  ----
  -c<canonicalization>    - r=relaxed [DEFAULT], s=simple, t=relaxed/simple, u=simple/relaxed
  -d<sdid>                - Signing Domain Identifier,if not provided it will be determined from the envelope originator/from header
  -i<auid>                - Agent User Identifier, usually the sender's email address (optional)
  -l                      - include body length tag (optional)
  -q                      - include query method tag
  -t                      - include a timestamp tag (optional)
  -x<expire_time>         - the expire time in seconds since epoch (optional, DEFAULT = current time + 604800)
  -y<selector>            - set RSA selector (DEFAULT: default)
  -Y<selector>            - set Ed25519 selector (DEFAULT: default)
  -z<hash>                - set signature type (1=sha1, 2=sha256, 3=both, 4=ed25519, 5=hybrid)
*/

int dkim_sign(const char *rsakeyfile, const char *ecckeyfile, const char *fnin, const char *fnout)
{
  int child;
  int wstat;
  char *(args[17]);
  int i = 0;

  args[i] = "qmail-dkim";
  ++i;
  args[i] = "-s";
  ++i;
  args[i] = "-q";
  ++i;
  if (sdid.len > 3) {
    args[i] = sdid.s;
    ++i;
  }
  if (selector.len > 3) {
    args[i] = selector.s;
    ++i;
  }
  if (selectore.len > 3) {
    args[i] = selectore.s;
    ++i;
  }
  if (auid.len > 3) {
    args[i] = auid.s;
    ++i;
  }
  if (expire.len > 3) {
    args[i] = expire.s;
    ++i;
  }
  if (canon.len > 2) {
    args[i] = canon.s;
    ++i;
  }
  if (hash.len > 2) {
    args[i] = hash.s;
    ++i;
  }
  if (length.len > 2) {
    args[i] = length.s;
    ++i;
  }
  args[i] = fnin;
  ++i;
  args[i] = rsakeyfile;
  ++i;
  args[i] = fnout;
  ++i;
  if (str_len(ecckeyfile) > 3) {
    args[i] = ecckeyfile;
    ++i;
  }
  args[i] = 0;

  if (!(child = vfork())) {
    pathexec(args);
    if (errno) _exit(111);
    _exit(100);
  }

  wait_pid(&wstat, child);
  if (wait_crashed(wstat)) return 1;

  switch (wait_exitcode(wstat)) {
    case 1:  return 1;
    default: return 0;
  }
}

int qmail_remote(char **qargs, int fd)
{
  int child;
  int wstat;
  char *(args[5]);

  args[0] = "qmail-remote";
  args[1] = qargs[1];
  args[2] = qargs[2];
  args[3] = qargs[3];
  args[4] = 0;

  if (!(child = vfork())) {
    if (fd) {
      if (fd_move(0, fd) == -1) _exit(111);
      if (fd_copy(2, 1) == -1) _exit(111);
    }
    pathexec(args);
    if (errno) _exit(111);
    _exit(100);
  }

  wait_pid(&wstat, child);
  if (wait_crashed(wstat)) return 1;

  switch (wait_exitcode(wstat)) {
    case 111: return 1;
    default:  return 0;
  }
}

void dkim_setup()
{
  int c, i, j, k, l;
  char *opt, *pos;

  /* defaults: selector=default, IETF format, q=dns/txt, z=2, c=r */

  if (!stralloc_copys(&sdid, "-d")) temp_nomem();
  if (!stralloc_cat(&sdid, &sender)) temp_nomem();
  if (!stralloc_0(&sdid)) temp_nomem();
  if (!stralloc_copys(&selector, "-ydefault")) temp_nomem();
  if (!stralloc_0(&selector)) temp_nomem();
  if (!stralloc_copys(&selectore, "-Yeddy")) temp_nomem();
  if (!stralloc_0(&selectore)) temp_nomem();
  if (!stralloc_copys(&canon, "-cr")) temp_nomem();
  if (!stralloc_0(&canon)) temp_nomem();
  if (!stralloc_copys(&hash, "-z2")) temp_nomem();
  if (!stralloc_0(&hash)) temp_nomem();

  /* domain:selector,selectore|sdid|[auid|~]|expire|c:z:l; c=[r|s|t|u], z=[1,2,3,4,5], l=l */

  if (dkimparams && *dkimparams) {
    i = str_chr(dkimparams, '|');
    pos = dkimparams + i;
    if (*pos == '|' || *pos == '\0') {  // selector
      dkimparams[i] = '\0';
      c = str_chr(dkimparams, ',');  // selectore=eddy
      if (dkimparams[c] == ',') {
        dkimparams[c] = '\0';
        if (str_len(dkimparams + c + 1)) {
          if (!stralloc_copys(&selectore, "-Y")) temp_nomem();
          if (!stralloc_cats(&selectore, dkimparams + c + 1)) temp_nomem();
          if (!stralloc_0(&selectore)) temp_nomem();
        }
      } else if (str_len(dkimparams)) {  // selector=default
        if (!stralloc_copys(&selector, "-y")) temp_nomem();
        if (!stralloc_cats(&selector, dkimparams)) temp_nomem();
        if (!stralloc_0(&selector)) temp_nomem();
      }

      j = str_chr(dkimparams + i + 1, '|');
      pos = dkimparams + i + j + 1;
      if (*pos == '|' || *pos == '\0') {  // sdid; domain in DKIM header
        dkimparams[i + j + 1] = '\0';
        if (!stralloc_copys(&sdid, "-d")) temp_nomem();
        if (!stralloc_cats(&sdid, dkimparams + i + 1)) temp_nomem();
        if (!stralloc_0(&sdid)) temp_nomem();

        k = str_chr(dkimparams + i + j + 2, '|');
        pos = dkimparams + i + j + k + 2;
        if (*pos == '|' || *pos == '\0') {  // auid = identifier
          dkimparams[i + j + k + 2] = '\0';
          if (!stralloc_copys(&auid, "-i")) temp_nomem();
          if (dkimparams[i + j + 2] == '~') {
            if (!stralloc_cat(&auid, &originator)) temp_nomem();
          } else if (!stralloc_cats(&auid, dkimparams + i + j + 2))
            temp_nomem();

          if (!stralloc_0(&auid)) temp_nomem();

          l = str_chr(dkimparams + i + j + k + 3, '|');
          pos = dkimparams + i + j + k + l + 3;
          if (*pos == '|' || *pos == '\0') {  // expire after n secs
            dkimparams[i + j + k + l + 3] = '\0';
            if (!stralloc_copys(&expire, "-x")) temp_nomem();
            if (!stralloc_cats(&expire, dkimparams + i + j + k + 3)) temp_nomem();
            if (!stralloc_0(&expire)) temp_nomem();

            /* Options to follow */

            opt = dkimparams + i + j + k + l + 4;
            if (*opt == '\0') return;
            if (*opt != ':') {
              if (!stralloc_copys(&canon, "-c")) temp_nomem();  // canonicalization
              if (!stralloc_catb(&canon, opt, 1)) temp_nomem();
              if (!stralloc_0(&canon)) temp_nomem();
              ++opt;
              if (*opt == '\0') return;  // next colon
            }
            if (*opt != ':' || *opt == '\0') return;
            if (*opt == ':') ++opt;
            if (*opt != ':') {
              if (!stralloc_copys(&hash, "-z")) temp_nomem();  // hash
              if (!stralloc_catb(&hash, opt, 1)) temp_nomem();
              if (!stralloc_0(&hash)) temp_nomem();
              ++opt;
              if (*opt == '\0') return;  // next colon
            }
            if (*opt != ':' || *opt == '\0') return;
            if (*opt == ':') ++opt;
            if (*opt != ':' && *opt == 'l') {
              if (!stralloc_copys(&length, "-l")) temp_nomem();  // length
              if (!stralloc_0(&length)) temp_nomem();
            }
          }
        }
      }
    }
  }

  return;
}

int main(int argc, char **args)
{
  int i;
  int fdin = 0;  // initial read from FD 0
  int nkey = 0;
  char *(qargs[4]);
  struct stat st;

  qargs[0] = args[0];
  qargs[1] = args[1];  // host
  qargs[2] = args[2];  // originator
  qargs[3] = args[3];  // recipient

  umask(033);
  sig_pipeignore();
  if (argc < 4) perm_usage();
  if (chdir(auto_qmail) == -1) temp_chdir();

  if (str_len(args[2]) > 2) {
    i = str_chr(args[2], '@');
    if (*(args[2] + i) == '@')
      if (!stralloc_copys(&senddomain, args[2] + i + 1)) temp_nomem();
  }
  if (!stralloc_0(&senddomain)) temp_nomem();
  if (!stralloc_copys(&originator, args[2])) temp_nomem();

  if (!get_controls()) {
    qmail_remote(qargs, fdin);
    _exit(0);
  }

  dkim_setup();  // sender is evaluated from originator (senddomain)

  /* Setup keys: they are composed from selector */

  case_lowerb(sender.s, sender.len);  // needs to be lowercase
  if (!stralloc_copys(&rsakey, DOMAINKEYS)) temp_nomem();
  if (!stralloc_cats(&rsakey, sender.s)) temp_nomem();
  if (!stralloc_cats(&rsakey, "/")) temp_nomem();

  if (!stralloc_copys(&ecckey, DOMAINKEYS)) temp_nomem();
  if (!stralloc_cats(&ecckey, sender.s)) temp_nomem();
  if (!stralloc_cats(&ecckey, "/")) temp_nomem();

  /* RSA key common for SHA1 and SHA256: rsakeyfile -> selector */

  if (!stralloc_cats(&rsakey, selector.s + 2)) temp_nomem();  // -y prepended
  if (!stralloc_0(&rsakey)) temp_nomem();
  if (stat(rsakey.s, &st) != -1)
    if (open_read(rsakey.s) > 0) ++nkey;

  /* ECC key follows: ecckeyfile -> (,)selector2 */

  if (!stralloc_cats(&ecckey, selectore.s + 2)) temp_nomem();  // -Y prepended
  if (!stralloc_0(&ecckey)) temp_nomem();
  if (stat(ecckey.s, &st) != -1)
    if (open_read(ecckey.s) > 0) ++nkey;

  /* We got keys - go for staging */

  if (nkey) {  // otherwise no key exists; why bother
    dkim_stage();
    if (!dkim_sign(rsakey.s, ecckey.s, fndkin.s, fndkout.s)) {
      fdin = open_read(fndkout.s);
      if (fdin == -1) die_read();
    } else {
      fdin = open_read(fndkin.s);  // DKIM key failed to sign
      if (fdin == -1) die_read();
    }
  } else {
    temp_nosignkey();
  }

  qmail_remote(qargs, fdin);  // closes fdin
  if (nkey) dkim_unlink();

  _exit(0);
}