// file      : openssl/agent/pkcs11/agent.cxx -*- C++ -*-
// copyright : Copyright (c) 2014-2018 Code Synthesis Ltd
// license   : MIT; see accompanying LICENSE file

#include <sys/un.h>     // sockaddr_un
#include <sys/types.h>
#include <sys/socket.h>

#include <signal.h>  // kill(), sigaction(), sigemptyset(), SIG*
#include <unistd.h>  // fork(), getpid(), dup2(), setsid()
#include <termios.h> // tcgetattr(), tcsetattr()

#include <csignal>  // sig_atomic_t
#include <cstring>  // strlen(), strcpy(), memset(), memcmp()
#include <iostream>

#include <iostream> // cout

#include <libbutl/pager.mxx>

#include <openssl/protocol.hxx>
#include <openssl/diagnostics.hxx>

#include <openssl/agent/pkcs11/url.hxx>
#include <openssl/agent/pkcs11/options.hxx>
#include <openssl/agent/pkcs11/private-key.hxx>

namespace openssl
{
  namespace agent
  {
    namespace pkcs11
    {
      using namespace std;
      using namespace butl;

      static void
      pin_prompt (const string& prompt, char*, size_t);

      // The agent daemon top-level function.
      //
      static int
      daemon (const options&,
              identity&&,
              access&&,
              auto_fd&& sock,
              char* pin) noexcept;

      static int
      main (int argc, char* argv[])
      try
      {
        cli::argv_scanner scan (argc, argv);

        // Parse options.
        //
        options ops;
        ops.parse (scan);

        // Version.
        //
        if (ops.version ())
        {
          cout << "openssl-agent-pkcs11 " << OPENSSL_AGENT_VERSION_ID << endl
               << "libbutl " << LIBBUTL_VERSION_ID << endl
               << "Copyright (c) 2014-2018 Code Synthesis Ltd" << endl
               << "This is free software released under the MIT license."
               << endl;

          return 0;
        }

        // Help.
        //
        if (ops.help ())
        {
          pager p ("openssl-agent-pkcs11 help", false);
          print_openssl_agent_pkcs11_usage (p.stream ());

          // If the pager failed, assume it has issued some diagnostics.
          //
          return p.wait () ? 0 : 1;
        }

        // Parse arguments.
        //
        identity key_identity;
        access   key_access;

        try
        {
          string u (scan.more () ? scan.next () : "");
          if (u.empty ())
            fail << "private key URL argument expected";

          url key_url (u);

          key_identity = identity (key_url);
          key_access   = access   (key_url);
        }
        catch (const invalid_argument& e)
        {
          fail << "invalid PKCS#11 URL: " << e;
        }

        try
        {
          // Prompt the user for a token PIN unless it is already specified as
          // an access attribute in the URL.
          //
          char pin[65]; // NULL character-terminated PIN.

          if (!key_access.pin_value)
            pin_prompt ("Enter PIN for PKCS#11", pin, sizeof (pin));

          // Open the agent communication socket and fail if the path is
          // already taken. We could probably deal with the faulty case
          // ourselves but let's not complicate things for now.
          //
          auto_fd sock (socket (AF_UNIX, SOCK_STREAM, 0));

          if (sock.get () == -1)
            throw_system_error (errno);

          path sock_path (path::temp_path ("openssl-agent-pkcs11"));

          struct sockaddr_un addr;
          memset (&addr, 0, sizeof (addr));
          addr.sun_family = AF_UNIX;

          if (sock_path.string ().size () >= sizeof (addr.sun_path))
            throw_generic_error (ENAMETOOLONG);

          strcpy (addr.sun_path, sock_path.string ().c_str ());

          // Creates the socket filesystem entry.
          //
          if (bind (sock.get (),
                    reinterpret_cast<sockaddr*> (&addr),
                    sizeof (addr)) == -1)
            throw_system_error (errno);

          auto_rmfile rm (sock_path);

          if (listen (sock.get (), 20 /* backlog */) == -1)
            throw_system_error (errno);

          // Fork off the agent daemon.
          //
          pid_t pid (fork ());

          if (pid == -1)
            throw_system_error (errno);

          // Run the agent daemon in the child process.
          //
          if (pid == 0)
            return daemon (ops,
                           move (key_identity),
                           move (key_access),
                           move (sock),
                           !key_access.pin_value ? pin : nullptr);

          // Erase the no longer needed PIN.
          //
          mem_clear (pin, sizeof (pin));

          // The child now is in charge for the socket filesystem entry.
          //
          rm.cancel ();

          sock.close ();

          // Send the init request to the agent to make sure that it is alive
          // and has successfully initialized the key. Issue the received from
          // the agent diagnostics and exit with non-zero status if the
          // initialization failed. Note that the daemon will also exit right
          // after the response in this case.
          //
          {
            sock = connect (sock_path);

            ifdstream is (fddup (sock.get ()));
            ofdstream os (move  (sock));

            os << request ("init");
            os.close ();

            response r;
            is >> r;

            if (r.status != 0)
            {
              text << r.error;
              return r.status;
            }
          }

          // Let's mimic the 'ssh-agent -s' command output.
          //
          cout << "OPENSSL_AGENT_PKCS11_SOCK=" << sock_path << "; "
               << "export OPENSSL_AGENT_PKCS11_SOCK;\n"
               << "OPENSSL_AGENT_PKCS11_PID=" << pid << "; "
               << "export OPENSSL_AGENT_PKCS11_PID;\n"
               << "echo Agent pid " << pid << endl;

          return 0;
        }
        catch (const io_error&)
        {
          fail << "unable to communicate with daemon";
        }
        catch (const system_error& e)
        {
          fail << "unable to start daemon: " << e;
        }

        return 0;
      }
      catch (const failed&)
      {
        return 1; // Diagnostics has already been issued.
      }
      catch (const cli::exception& e)
      {
        error << e;
        return 1;
      }

      // Prompt for the PIN.
      //
      // We would be happy to use getpass() functions but it is mentioned as
      // obsolete in the man page and is likely to be removed from the glibc
      // future versions. Thus, we will provide our own implementation of the
      // function that is inspired by the openssh implementation.
      //
      // Note: _NSIG is Linux-specic.
      //
      static volatile sig_atomic_t received_signals[_NSIG];

      extern "C" void
      handle_signal (int sig)
      {
        received_signals[sig] = 1;
      }

      static void
      pin_prompt (const string& prompt, char* buf, size_t max_len)
      {
        // Note that openssh tries to open /dev/tty for read/write access but
        // falls back to stdin/stdout on failure. Let's skip the fallback for
        // simplicity.
        //
        try
        {
          auto_fd ifd  (fdopen ("/dev/tty", fdopen_mode::in));
          ofdstream os (fdopen ("/dev/tty", fdopen_mode::out));

          // Disable the echo.
          //
          termios term_attrs;
          if (tcgetattr (os.fd (), &term_attrs) == -1)
            throw_system_error (errno);

          termios na (term_attrs);
          na.c_lflag &= ~ECHO;

          if (tcsetattr (os.fd (), TCSAFLUSH, &na) == -1)
            throw_system_error (errno);

          // Catch signals that would otherwise cause the user to end up with
          // echo turned off in the shell.
          //
          for (size_t i (0); i < _NSIG; ++i)
            received_signals[i] = 0;

          using sigaction_type = struct sigaction;

          sigaction_type sa;
          memset (&sa, 0, sizeof (sa));
          sigemptyset (&sa.sa_mask);
          sa.sa_handler = handle_signal;

          // Save the current signal handlers while setting up the new ones.
          //
          vector<pair<int, sigaction_type>> sig_handlers;

          const vector<int> sigs ({
              SIGALRM,
              SIGHUP,
              SIGINT,
              SIGPIPE,
              SIGQUIT,
              SIGTERM,
              SIGTSTP,
              SIGTTIN,
              SIGTTOU});

          for (int s: sigs)
          {
            sig_handlers.push_back (make_pair (s, sigaction_type ()));

            if (sigaction (s, &sa, &sig_handlers.back ().second) == -1)
              throw_system_error (errno);
          }

          // Restore the echo state and signal handlers and resend the
          // received signals.
          //
          auto restore = [&os, &term_attrs, &sig_handlers] ()
          {
            if (tcsetattr (os.fd (), TCSAFLUSH, &term_attrs) == -1)
              throw_system_error (errno);

            for (const pair<int, sigaction_type>& sa: sig_handlers)
            {
              if (sigaction (sa.first, &sa.second, nullptr) == -1)
                throw_system_error (errno);
            }

            pid_t pid (getpid ());
            for (int i (0); i < _NSIG; i++)
            {
              if (received_signals[i] != 0 && kill (pid, i) == -1)
                throw_system_error (errno);
            }
          };

          auto rst (make_exception_guard ([&restore] ()
                                          {
                                            try
                                            {
                                              restore ();
                                            }
                                            catch (const system_error&)
                                            {
                                              // Not much we can do here.
                                              //
                                            }
                                          }));

          // Print prompt.
          //
          os << prompt << ": " << flush;

          // Read the PIN.
          //
          // Reserve space for the terminating NULL character.
          //
          size_t n (max_len - 1);
          size_t i (0);
          for (; i < n; ++i)
          {
            ssize_t r (read (ifd.get (), buf + i, 1));
            if (r == -1)
              throw_system_error (errno);

            if (r == 0 || buf[i] == '\n')
              break;
          }

          restore ();

          os << endl;
          os.close ();

          if (i == n)
            fail << "PIN is too long";

          buf[i] = '\0';
        }
        catch (const system_error& e)
        {
          fail << "failed to read PIN: " << e;
        }
      }

      // Terminate the daemon.
      //
      static volatile sig_atomic_t exit_requested (0); // Can be used as bool.

      extern "C" void
      handle_term (int /* sig */)
      {
        exit_requested = 1;

        // Now accept() will wakeup, set errno to EINTR, and return with -1.
        //
      }

      // Run the daemon.
      //
      static int
      daemon (const options& ops,
              identity&& ki,
              access&& ka,
              auto_fd&& sc,
              char* pin) noexcept
      try
      {
        // Demonize ourselves.
        //
        {
          if (setsid () == -1)
            throw_system_error (errno);

          path ("/").current_directory ();

          auto_fd nd (fdnull ());

          auto redir = [&nd] (int sd)
          {
            if (dup2 (nd.get (), sd) == -1)
              throw_system_error (errno);
          };

          redir (0);
          redir (1);
          redir (2);
        }

        auto_fd  sock         (move (sc));
        identity key_identity (move (ki));
        access   key_access   (move (ka));

        struct sigaction sa;
        memset (&sa, 0, sizeof (sa));
        sigemptyset (&sa.sa_mask);
        sa.sa_handler = handle_term;

        if (sigaction (SIGTERM, &sa, nullptr) == -1)
          throw_system_error (errno);

        // It would be natural to create the key at the time of receiving the
        // init request. However, we need to erase the PIN as soon as
        // possible. Thus, we will create it now. In case of failure we will
        // save the error to use for the upcoming init request response.
        //
        private_key key;

        // If present, indicates that we still expect the init request. Will
        // set it to nullopt after init request processing.
        //
        optional<string> init_error;

        try
        {
          optional<simulate_outcome> s;
          if (ops.simulate_specified ())
            s = ops.simulate ();

          key = private_key (key_identity, key_access, pin, s);
          init_error = "";
        }
        catch (const invalid_argument& e)
        {
          init_error = string ("error: ") + e.what ();
        }
        catch (const runtime_error& e)
        {
          init_error = string ("error: ") + e.what ();
        }

        // Erase the no longer needed PIN.
        //
        if (pin != nullptr)
          mem_clear (pin, strlen (pin));

        // Run the request processing loop.
        //
        while (true)
        {
          // Accept the client connection.
          //
          auto_fd fd (
            accept (sock.get (), nullptr /* addr */, nullptr /* addrlen */));

          if (exit_requested)
            break;

          if (fd.get () == -1)
          {
            if (errno == EINTR || errno == ECONNABORTED || errno == EPROTO)
              continue;

            throw_system_error (errno);
          }

          // If we fail to duplicate the file descriptor, then we are probably
          // out of resources. Anyway we can't do much about it and will bail
          // out.
          //
          ifdstream is (fddup (fd.get ()));

          try
          {
            ofdstream os (move (fd));

            // Respond to the client until he sends valid requests.
            //
            while (is.peek () != ifdstream::traits_type::eof ())
            {
              request rq;
              is >> rq;

              // The init request must be the first and only one.
              //
              if (rq.cmd == "init")
              {
                if (!init_error)
                {
                  os << response ("error: already initialized");
                  break;
                }

                if (!init_error->empty ())
                {
                  os << response (*init_error);
                  os.close ();

                  throw runtime_error (*init_error);
                }

                init_error = nullopt; // We don't expect init anymore.
                os << response ();    // Success.
                continue;
              }

              if (init_error)
              {
                os << response ("error: not initialized");
                break;
              }

              // Process the sign request.
              //
              if (rq.cmd == "sign")
              {
                // Parse and validate the request arguments (key URL and
                // simulate outcome).
                //
                if (rq.args.size () != 2)
                {
                  os << response ("error: invalid args");
                  break;
                }

                identity k;

                try
                {
                  k = identity (url (rq.args[0]));
                }
                catch (const invalid_argument& e)
                {
                  os << response (
                    string ("error: invalid private key URL: ") + e.what ());

                  break;
                }

                optional<simulate_outcome> sim;

                if (!rq.args[1].empty ())
                try
                {
                  sim = to_simulate_outcome (rq.args[1]);
                }
                catch (const invalid_argument& e)
                {
                  os << response (string ("error: ") + e.what ());
                  break;
                }

                if (k != key_identity)
                {
                  os << response ("error: private key doesn't match");
                  break;
                }

                try
                {
                  os << response (key.sign (rq.input, sim));
                }
                catch (const runtime_error& e)
                {
                  os << response (string ("error: ") + e.what ());
                  break;
                }
              }
            }

            // Close the client connection.
            //
            is.close ();
            os.close ();
          }
          catch (const io_error&)
          {
            // The client is presumably dead.
            //
          }
        }

        return 0;
      }
      catch (const std::exception&)
      {
        // There is no way we can complain.
        //
        return 1;
      }
    }
  }
}

int
main (int argc, char* argv[])
{
  return openssl::agent::pkcs11::main (argc, argv);
}