#include <errno.h>
#include <linux/limits.h>
#include <netdb.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/types.h>

#include "includes.h"

#include "hostfile.h"   /* Needs to be included before auth.h */
#include "auth.h"
#include "kex.h"
#include "log.h"
#include "misc.h"
#include "monitor.h"
#include "ssh-gss.h"    /* Needs to be included before monitor_wrap.h */
#include "monitor_wrap.h"
#include "pathnames.h"
#include "servconf.h"
#include "sshbuf.h"

#define MAX_LISTEN_STREAMS      (16)
#define MAX_LISTEN_STREAM_LEN   (NI_MAXHOST + NI_MAXSERV + sizeof("ListenAddress=[:]") + 1)
typedef char listen_stream_set[MAX_LISTEN_STREAMS][MAX_LISTEN_STREAM_LEN];

/* Global variables required for sshd config parsing. */
ServerOptions options = {};
struct sshbuf *cfg = NULL;
struct include_list includes = TAILQ_HEAD_INITIALIZER(includes);

/* Other global variables that are required for this to build, because of their
 * use throughout the codebase. We do NOT use these variables for the
 * generator. */
Authctxt *the_authctxt = NULL;
int privsep_is_preauth = 1;
int use_privsep = -1;
struct monitor *pmonitor = NULL;
struct ssh *the_active_state = NULL;
struct sshauthopt *auth_opts = NULL;
struct sshbuf *loginmsg = NULL;

static int listen_stream_set_append(listen_stream_set set, const char *listen_stream) {
        size_t n;

        if (!set)
                return -EINVAL;

        n = strnlen(listen_stream, MAX_LISTEN_STREAM_LEN);
        if (n == MAX_LISTEN_STREAM_LEN)
                return -EINVAL;

        for (int i = 0; i < MAX_LISTEN_STREAMS; i++) {
                if (strcmp(set[i], listen_stream) == 0)
                        return 0;

                if (strnlen(set[i], MAX_LISTEN_STREAM_LEN) > 0)
                        continue;

                memcpy(set[i], listen_stream, n);

                return 0;
        }

        return -E2BIG;
}

static int listen_stream_set_len(listen_stream_set set) {
        int r = 0;

        if (!set)
                return 0;

        for (int i = 0; i < MAX_LISTEN_STREAMS; i++) {
                if (strnlen(set[i], MAX_LISTEN_STREAM_LEN) > 0)
                        r++;
                else
                        break;
        }

        return r;
}

static char *path_append(const char *base, const char *append) {
        bool add_slash;
        size_t n = 0, len_base, len_append;
        char *path = NULL;

        len_base = strnlen(base, PATH_MAX);
        len_append = strnlen(append, PATH_MAX);
        add_slash = base[len_base - 1] != '/';

        path = calloc(len_base + len_append + (add_slash ? 2 : 1), sizeof(char));
        if (!path)
                return NULL;

        memcpy(path, base, len_base);
        n += len_base;

        if (add_slash)
                path[n++] = '/';

        memcpy(path + n, append, len_append);
        n += len_append;
        path[n] = '\0';

        return path;
}

static int fflush_and_check(FILE *f) {
        errno = 0;
        fflush(f);

        if (ferror(f))
                return errno > 0 ? -errno : -EIO;

        return 0;
}

static int write_systemd_socket_file(const char *destdir) {
        listen_stream_set listen_streams = {};
        int num_listen_streams;
        char *conf = NULL, *overridedir = NULL;
        FILE *f = NULL;
        int r;

        overridedir = path_append(destdir, "ssh.socket.d");
        if (!overridedir) {
                r = -ENOMEM;
                goto out;
        }

        if (mkdir(overridedir, 0755) < 0 && errno != EEXIST) {
                r = -errno;
                goto out;
        }

        conf = path_append(overridedir, "addresses.conf");
        if (!conf) {
                r = -ENOMEM;
                goto out;
        }

        f = fopen(conf, "we");
        if (!f) {
                r = -errno;
                goto out;
        }

        fprintf(f,
                "# Automatically generated by sshd-socket-generator\n"
                "\n[Socket]\n"
                "ListenStream=\n");

        for (u_int i = 0; i < options.num_listen_addrs; i++) {
                for (struct addrinfo *ai = options.listen_addrs[i].addrs; ai; ai = ai->ai_next) {
                        char addr[NI_MAXHOST] = {}, port[NI_MAXSERV] = {},
                             listen_stream[MAX_LISTEN_STREAM_LEN] = {};

                        r = getnameinfo(ai->ai_addr, ai->ai_addrlen,
                                        addr, sizeof(addr),
                                        port, sizeof(port),
                                        NI_NUMERICHOST|NI_NUMERICSERV);
                        if (r != 0) {
                                fprintf(stderr, "%s\n", gai_strerror(r));
                                r = r == EAI_SYSTEM ? -errno : -EINVAL;
                                goto out;
                        }

                        if (strcmp(addr, "0.0.0.0") == 0 || strcmp(addr, "::") == 0) {
                                /* If ListenAddress is 0.0.0.0 or ::, only
                                 * write the port in ListenStream=. */
                                snprintf(listen_stream,
                                        MAX_LISTEN_STREAM_LEN,
                                        "ListenStream=%s",
                                        port);
                        } else
                                snprintf(listen_stream,
                                        MAX_LISTEN_STREAM_LEN,
                                        "ListenStream=%s%s%s:%s",
                                        ai->ai_family == AF_INET6 ? "[" : "",
                                        addr,
                                        ai->ai_family == AF_INET6 ? "]" : "",
                                        port);

                        r = listen_stream_set_append(listen_streams, listen_stream);
                        if (r < 0)
                                goto out;
                }
        }

        num_listen_streams = listen_stream_set_len(listen_streams);

        if (num_listen_streams <= 0) {
                /* We didn't generate anything useful, so clean up and leave
                 * ssh.socket as-is. */
                r = -ENODATA;
                goto out;
        }

        if (num_listen_streams == 1 && strcmp(listen_streams[0], "ListenStream=22") == 0) {
                /* This is the default already specified in ssh.socket. No need
                 * to write the override. */
                r = -ENODATA;
                goto out;
        }

        for (int i = 0; i < num_listen_streams; i++)
                fprintf(f, "%s\n", listen_streams[i]);

        r = fflush_and_check(f);
        if (r < 0)
                goto out;

out:
        if (f)
                fclose(f);

        if (r < 0) {
                (void) remove(conf);
                (void) remove(overridedir);
        }

        free(overridedir);
        free(conf);

        return r;
}

static int parse_sshd_config_options() {
        struct connection_info *connection_info;

        cfg = sshbuf_new();
        if (!cfg)
                return -ENOMEM;

	initialize_server_options(&options);
        load_server_config(_PATH_SERVER_CONFIG_FILE, cfg);
        parse_server_config(&options, _PATH_SERVER_CONFIG_FILE, cfg, &includes, NULL, 0);
        fill_default_server_options(&options);

        connection_info = get_connection_info(NULL, 0, 0);
        connection_info->test = 1;

        parse_server_match_config(&options, &includes, connection_info);

        return 0;
}

int main(int argc, char **argv) {
        const char *destdir = NULL;
        int r;

        if (argc < 2) {
                fprintf(stderr, "Expected at least one argument.\n");

                return EXIT_FAILURE;
        }

        destdir = argv[1];

        r = parse_sshd_config_options();
        if (r < 0) {
                fprintf(stderr, "Faild to parse sshd config: %s\n", strerror(-r));

                return EXIT_FAILURE;
        }

        if (options.num_listen_addrs <= 0) {
                /* No listen addresses configured? Don't generate anything. */
                fprintf(stderr, "No listen addresses configured. Will not generate anything.\n");

                return EXIT_SUCCESS;
        }

        r = write_systemd_socket_file(destdir);
        if (r == -ENODATA) {
                fprintf(stderr, "No custom listen addresses configured. Will not generated anything.\n");

                return EXIT_SUCCESS;
        }
        if (r < 0) {
                fprintf(stderr, "Failed to generate ssh.socket: %s\n", strerror(-r));

                return EXIT_FAILURE;
        }

        return EXIT_SUCCESS;
}