/*
 * Copyright (c) 2022 Genexis B.V.
 *
 * This program and the accompanying materials are made available under the
 * terms of the Eclipse Public License 2.0 which is available at
 * https://www.eclipse.org/legal/epl-2.0/
 *
 * SPDX-License-Identifier: EPL-2.0
 *
 * Contributors:
 *   Erik Karlsson - initial implementation
 */

#define _GNU_SOURCE
#include <stdio.h>
#include <string.h>
#include <shadow.h>
#include <crypt.h>
#include <stdlib.h>
#include <arpa/inet.h>
#include <mosquitto.h>
#include <mosquitto_broker.h>
#include <mosquitto_plugin.h>

#ifdef ENABLE_PAM_SUPPORT
#include <security/pam_appl.h>
#endif

#define MAX_USERS 256
#define MAX_SUBNETS_PER_USER 32

typedef struct {
	union {
		uint32_t ipv4_network;
		uint8_t ipv6_network[16];
	};
	union {
		uint32_t ipv4_netmask;
		uint8_t ipv6_netmask[16];
	};
	int is_ipv6;
} subnet_t;

typedef struct {
	char username[64];
	subnet_t allow_subnets[MAX_SUBNETS_PER_USER];
	int allow_count;
	subnet_t deny_subnets[MAX_SUBNETS_PER_USER];
	int deny_count;
} user_acl_t;

typedef struct {
	user_acl_t users[MAX_USERS];
	int user_count;
	mosquitto_plugin_id_t *identifier;
	char *config_file;
} plugin_data_t;

/* Parse CIDR notation for IPv4 or IPv6 (e.g., "192.168.1.0/24" or "2001:db8::/32") */
static int parse_subnet(const char *cidr, subnet_t *subnet)
{
	char ip_str[128];
	char *slash;
	int prefix_len;
	struct in_addr addr4;
	struct in6_addr addr6;

	strncpy(ip_str, cidr, sizeof(ip_str) - 1);
	ip_str[sizeof(ip_str) - 1] = '\0';

	slash = strchr(ip_str, '/');
	if (slash != NULL) {
		*slash = '\0';
		prefix_len = atoi(slash + 1);
	}

	/* Try IPv4 first */
	if (inet_pton(AF_INET, ip_str, &addr4) == 1) {
		subnet->is_ipv6 = 0;
		if (slash == NULL)
			prefix_len = 32;
		if (prefix_len < 0 || prefix_len > 32)
			return -1;

		subnet->ipv4_network = ntohl(addr4.s_addr);
		subnet->ipv4_netmask = prefix_len == 0 ? 0 : (~0U << (32 - prefix_len));
		subnet->ipv4_network &= subnet->ipv4_netmask;
		return 0;
	}

	/* Try IPv6 */
	if (inet_pton(AF_INET6, ip_str, &addr6) == 1) {
		subnet->is_ipv6 = 1;
		if (slash == NULL)
			prefix_len = 128;
		if (prefix_len < 0 || prefix_len > 128)
			return -1;

		/* Copy network address */
		memcpy(subnet->ipv6_network, addr6.s6_addr, 16);

		/* Generate netmask */
		memset(subnet->ipv6_netmask, 0, 16);
		for (int i = 0; i < prefix_len / 8; i++)
			subnet->ipv6_netmask[i] = 0xff;
		if (prefix_len % 8)
			subnet->ipv6_netmask[prefix_len / 8] = ~((1 << (8 - (prefix_len % 8))) - 1);

		/* Apply netmask to network address */
		for (int i = 0; i < 16; i++)
			subnet->ipv6_network[i] &= subnet->ipv6_netmask[i];

		return 0;
	}

	return -1;
}

/* Check if IPv4 address is in subnet */
static int ipv4_in_subnet(uint32_t ip, const subnet_t *subnet)
{
	if (subnet->is_ipv6)
		return 0;
	return (ip & subnet->ipv4_netmask) == subnet->ipv4_network;
}

/* Check if IPv6 address is in subnet */
static int ipv6_in_subnet(const uint8_t *ip, const subnet_t *subnet)
{
	if (!subnet->is_ipv6)
		return 0;
	for (int i = 0; i < 16; i++) {
		if ((ip[i] & subnet->ipv6_netmask[i]) != subnet->ipv6_network[i])
			return 0;
	}
	return 1;
}

/* Check if IP is in any subnet in the list */
static int ip_in_subnet_list(const char *client_address, const subnet_t *subnets, int count)
{
	struct in_addr addr4;
	struct in6_addr addr6;
	uint32_t ipv4;

	/* Try IPv4 */
	if (inet_pton(AF_INET, client_address, &addr4) == 1) {
		ipv4 = ntohl(addr4.s_addr);
		for (int i = 0; i < count; i++) {
			if (ipv4_in_subnet(ipv4, &subnets[i]))
				return 1;
		}
		return 0;
	}

	/* Try IPv6 */
	if (inet_pton(AF_INET6, client_address, &addr6) == 1) {
		for (int i = 0; i < count; i++) {
			if (ipv6_in_subnet(addr6.s6_addr, &subnets[i]))
				return 1;
		}
		return 0;
	}

	return 0;
}

/* Find or create user ACL entry */
static user_acl_t* find_or_create_user_acl(plugin_data_t *pdata, const char *username)
{
	user_acl_t *user;

	/* Find existing user */
	for (int i = 0; i < pdata->user_count; i++) {
		if (strcmp(pdata->users[i].username, username) == 0)
			return &pdata->users[i];
	}

	/* Create new user if not found */
	if (pdata->user_count >= MAX_USERS) {
		mosquitto_log_printf(MOSQ_LOG_ERR,
			"subnet_acl: Max users exceeded");
		return NULL;
	}

	user = &pdata->users[pdata->user_count];
	strncpy(user->username, username, sizeof(user->username) - 1);
	user->username[sizeof(user->username) - 1] = '\0';
	user->allow_count = 0;
	user->deny_count = 0;
	pdata->user_count++;

	return user;
}

/* Parse subnet ACL file with simplified format
 * Format:
 * # Comment lines
 * subnet allow <username> <cidr>
 * subnet deny <username> <cidr>
 */
static int load_subnet_acl_config(plugin_data_t *pdata, const char *config_file)
{
	FILE *fp;
	char line[512];
	int line_num = 0;

	/* Initialize user count */
	pdata->user_count = 0;

	/* Config file is optional - if not provided, no subnet filtering */
	if (config_file == NULL) {
		mosquitto_log_printf(MOSQ_LOG_INFO,
			"subnet_acl: No subnet ACL file specified, subnet filtering disabled");
		return 0;
	}

	/* If config file is specified but cannot be opened, this is a fatal error */
	fp = fopen(config_file, "r");
	if (fp == NULL) {
		mosquitto_log_printf(MOSQ_LOG_ERR,
			"subnet_acl: Failed to open subnet ACL file '%s'", config_file);
		return -1;
	}

	while (fgets(line, sizeof(line), fp) != NULL) {
		char *token, *saveptr;
		char *action, *username, *cidr;
		user_acl_t *user;
		subnet_t subnet;

		line_num++;

		/* Remove newline and comments */
		line[strcspn(line, "\r\n")] = '\0';
		char *comment = strchr(line, '#');
		if (comment)
			*comment = '\0';

		/* Trim leading whitespace */
		char *line_start = line;
		while (*line_start == ' ' || *line_start == '\t')
			line_start++;

		/* Skip empty lines */
		if (*line_start == '\0')
			continue;

		/* Parse: subnet allow|deny <username> <cidr> */
		token = strtok_r(line_start, " \t", &saveptr);
		if (token == NULL)
			continue;

		/* Must start with "subnet" */
		if (strcmp(token, "subnet") != 0) {
			mosquitto_log_printf(MOSQ_LOG_ERR,
				"subnet_acl: Invalid directive '%s' at line %d (expected 'subnet')",
				token, line_num);
			fclose(fp);
			return -1;
		}

		/* Get allow/deny */
		action = strtok_r(NULL, " \t", &saveptr);
		if (action == NULL) {
			mosquitto_log_printf(MOSQ_LOG_ERR,
				"subnet_acl: Missing allow/deny at line %d", line_num);
			fclose(fp);
			return -1;
		}

		if (strcmp(action, "allow") != 0 && strcmp(action, "deny") != 0) {
			mosquitto_log_printf(MOSQ_LOG_ERR,
				"subnet_acl: Invalid action '%s' at line %d (use 'allow' or 'deny')",
				action, line_num);
			fclose(fp);
			return -1;
		}

		/* Get username */
		username = strtok_r(NULL, " \t", &saveptr);
		if (username == NULL) {
			mosquitto_log_printf(MOSQ_LOG_ERR,
				"subnet_acl: Missing username at line %d", line_num);
			fclose(fp);
			return -1;
		}

		/* Get CIDR */
		cidr = strtok_r(NULL, " \t", &saveptr);
		if (cidr == NULL) {
			mosquitto_log_printf(MOSQ_LOG_ERR,
				"subnet_acl: Missing CIDR at line %d", line_num);
			fclose(fp);
			return -1;
		}

		/* Parse subnet */
		if (parse_subnet(cidr, &subnet) != 0) {
			mosquitto_log_printf(MOSQ_LOG_ERR,
				"subnet_acl: Invalid CIDR '%s' at line %d", cidr, line_num);
			fclose(fp);
			return -1;
		}

		/* Find or create user */
		user = find_or_create_user_acl(pdata, username);
		if (user == NULL) {
			mosquitto_log_printf(MOSQ_LOG_ERR,
				"subnet_acl: Max users (%d) exceeded at line %d", MAX_USERS, line_num);
			fclose(fp);
			return -1;
		}

		/* Add to appropriate list */
		if (strcmp(action, "allow") == 0) {
			if (user->allow_count >= MAX_SUBNETS_PER_USER) {
				mosquitto_log_printf(MOSQ_LOG_ERR,
					"subnet_acl: Max allow subnets (%d) exceeded for user '%s' at line %d",
					MAX_SUBNETS_PER_USER, user->username, line_num);
				fclose(fp);
				return -1;
			}
			user->allow_subnets[user->allow_count] = subnet;
			user->allow_count++;

			mosquitto_log_printf(MOSQ_LOG_DEBUG,
				"subnet_acl: User '%s' allow subnet %s",
				user->username, cidr);

		} else { /* deny */
			if (user->deny_count >= MAX_SUBNETS_PER_USER) {
				mosquitto_log_printf(MOSQ_LOG_ERR,
					"subnet_acl: Max deny subnets (%d) exceeded for user '%s' at line %d",
					MAX_SUBNETS_PER_USER, user->username, line_num);
				fclose(fp);
				return -1;
			}
			user->deny_subnets[user->deny_count] = subnet;
			user->deny_count++;

			mosquitto_log_printf(MOSQ_LOG_DEBUG,
				"subnet_acl: User '%s' deny subnet %s",
				user->username, cidr);
		}
	}

	fclose(fp);

	/* Log summary */
	for (int i = 0; i < pdata->user_count; i++) {
		user_acl_t *user = &pdata->users[i];
		if (user->allow_count > 0 || user->deny_count > 0) {
			mosquitto_log_printf(MOSQ_LOG_INFO,
				"subnet_acl: User '%s' has %d allow and %d deny subnet rules",
				user->username, user->allow_count, user->deny_count);
		}
	}

	mosquitto_log_printf(MOSQ_LOG_NOTICE,
		"subnet_acl: Loaded subnet restrictions for %d user(s)", pdata->user_count);

	return 0;
}

/* Find user ACL entry */
static const user_acl_t* find_user_acl(const plugin_data_t *pdata, const char *username)
{
	for (int i = 0; i < pdata->user_count; i++) {
		if (strcmp(pdata->users[i].username, username) == 0)
			return &pdata->users[i];
	}
	return NULL;
}

/* Check subnet access on authentication (connection time)
 * Returns: MOSQ_ERR_SUCCESS if allowed, MOSQ_ERR_AUTH if denied
 */
static int check_subnet_on_auth(plugin_data_t *pdata, struct mosquitto_evt_basic_auth *ed)
{
	const user_acl_t *user_acl;
	const char *client_address;

	/* Skip if no subnet config loaded */
	if (pdata == NULL || pdata->user_count == 0)
		return MOSQ_ERR_SUCCESS;

	/* Skip anonymous users */
	if (ed->username == NULL)
		return MOSQ_ERR_SUCCESS;

	/* Find user's subnet ACL */
	user_acl = find_user_acl(pdata, ed->username);

	/* If user not in config or has no subnet rules, allow */
	if (user_acl == NULL || (user_acl->allow_count == 0 && user_acl->deny_count == 0))
		return MOSQ_ERR_SUCCESS;

	/* Get client IP address */
	client_address = mosquitto_client_address(ed->client);
	if (client_address == NULL) {
		mosquitto_log_printf(MOSQ_LOG_WARNING,
			"subnet_acl: Could not get client address for user '%s', denying connection",
			ed->username);
		return MOSQ_ERR_AUTH;
	}

	/* Check deny list first - deny takes precedence */
	if (user_acl->deny_count > 0) {
		if (ip_in_subnet_list(client_address, user_acl->deny_subnets, user_acl->deny_count)) {
			mosquitto_log_printf(MOSQ_LOG_NOTICE,
				"subnet_acl: User '%s' from %s DENIED by deny rule",
				ed->username, client_address);
			return MOSQ_ERR_AUTH;
		}
	}

	/* If there are allow rules, IP must match one of them */
	if (user_acl->allow_count > 0) {
		if (ip_in_subnet_list(client_address, user_acl->allow_subnets, user_acl->allow_count)) {
			mosquitto_log_printf(MOSQ_LOG_DEBUG,
				"subnet_acl: User '%s' from %s allowed by allow rule",
				ed->username, client_address);
			return MOSQ_ERR_SUCCESS;
		} else {
			mosquitto_log_printf(MOSQ_LOG_NOTICE,
				"subnet_acl: User '%s' from %s DENIED (not in allowed subnets)",
				ed->username, client_address);
			return MOSQ_ERR_AUTH;
		}
	}

	/* No subnet rules for this user - allow */
	return MOSQ_ERR_SUCCESS;
}

#ifdef ENABLE_PAM_SUPPORT
static int pam_conversation(int num_msg, const struct pam_message **msg, struct pam_response **resp, void *appdata_ptr)
{
	int i;
	const char *pass = (const char *)appdata_ptr;

	*resp = calloc(num_msg, sizeof(struct pam_response));
	if (*resp == NULL) {
		mosquitto_log_printf(MOSQ_LOG_ERR, "pam failed to allocate buffer for validation");
		return PAM_BUF_ERR;
	}

	if (pass == NULL)
		return PAM_SUCCESS;

	for (i = 0; i < num_msg; ++i) {
		if (msg[i]->msg_style == PAM_PROMPT_ECHO_OFF) {
			(*resp)[i].resp = strdup(pass);
			if ((*resp)[i].resp == NULL) {
				for (int j = 0; j < i ; j++)
					free((*resp)[j].resp);

				free(*resp);
				*resp = NULL;
				mosquitto_log_printf(MOSQ_LOG_ERR, "pam failed in strdup");
				return PAM_BUF_ERR;
			}
		}
	}
	return PAM_SUCCESS;
}

static int process_pam_auth_callback(struct mosquitto_evt_basic_auth *ed)
{
	struct pam_conv conv;
	int retval;
	pam_handle_t *pamh = NULL;

	conv.conv = pam_conversation;
	conv.appdata_ptr = (void *)ed->password;

	retval = pam_start("mosquitto", ed->username, &conv, &pamh);
	if (retval != PAM_SUCCESS) {
		mosquitto_log_printf(MOSQ_LOG_ERR, "pam start failed: %s", pam_strerror(pamh, retval));
		return MOSQ_ERR_AUTH;
	}

	retval = pam_authenticate(pamh, 0);
	pam_end(pamh, retval);
	if (retval == PAM_SUCCESS) {
		mosquitto_log_printf(MOSQ_LOG_NOTICE, "pam user [%s] logged in", ed->username);
		return MOSQ_ERR_SUCCESS;
	}

	mosquitto_log_printf(MOSQ_LOG_NOTICE, "pam user [%s] failed authentication, err [%s]", ed->username, pam_strerror(pamh, retval));
	return MOSQ_ERR_AUTH;
}
#else
static int process_shadow_auth_callback(struct mosquitto_evt_basic_auth *ed)
{
	struct spwd spbuf, *sp = NULL;
	char buf[256];
	struct crypt_data data;
	char *hash;

	getspnam_r(ed->username, &spbuf, buf, sizeof(buf), &sp);

	if (sp == NULL || sp->sp_pwdp == NULL)
		return MOSQ_ERR_AUTH;

	/* Empty string as hash means password is not required */
	if (sp->sp_pwdp[0] == 0)
		return MOSQ_ERR_SUCCESS;

	if (ed->password == NULL)
		return MOSQ_ERR_AUTH;

	memset(&data, 0, sizeof(data));
	hash = crypt_r(ed->password, sp->sp_pwdp, &data);

	if (hash == NULL)
		return MOSQ_ERR_AUTH;

	if (strcmp(hash, sp->sp_pwdp) == 0)
		return MOSQ_ERR_SUCCESS;

	return MOSQ_ERR_AUTH;
}
#endif

static int basic_auth_callback(int event, void *event_data, void *userdata)
{
	struct mosquitto_evt_basic_auth *ed = event_data;
	plugin_data_t *pdata = userdata;
	int auth_result;

	/* Let other plugins or broker decide about anonymous login */
	if (ed->username == NULL)
		return MOSQ_ERR_PLUGIN_DEFER;

	/* First check username/password authentication */
#ifdef ENABLE_PAM_SUPPORT
	auth_result = process_pam_auth_callback(ed);
#else
	auth_result = process_shadow_auth_callback(ed);
#endif

	/* If authentication failed, reject immediately */
	if (auth_result != MOSQ_ERR_SUCCESS)
		return auth_result;

	/* Authentication succeeded, now check subnet restrictions */
	return check_subnet_on_auth(pdata, ed);
}

static int reload_callback(int event, void *event_data, void *userdata)
{
	plugin_data_t *pdata = userdata;

	if (pdata == NULL)
		return MOSQ_ERR_SUCCESS;

	mosquitto_log_printf(MOSQ_LOG_NOTICE,
		"subnet_acl: Reloading subnet ACL configuration from '%s'",
		pdata->config_file ? pdata->config_file : "(none)");

	/* Reload subnet ACL configuration */
	if (load_subnet_acl_config(pdata, pdata->config_file) != 0) {
		mosquitto_log_printf(MOSQ_LOG_ERR,
			"subnet_acl: Failed to reload subnet ACL configuration, keeping old config");
		return MOSQ_ERR_UNKNOWN;
	}

	mosquitto_log_printf(MOSQ_LOG_NOTICE,
		"subnet_acl: Reload complete, now tracking %d user(s)", pdata->user_count);

	return MOSQ_ERR_SUCCESS;
}

int mosquitto_plugin_version(int supported_version_count,
			     const int *supported_versions)
{
	return 5;
}

int mosquitto_plugin_init(mosquitto_plugin_id_t *identifier,
			  void **user_data,
			  struct mosquitto_opt *opts, int opt_count)
{
	plugin_data_t *pdata;
	const char *config_file = NULL;
	int rc;

	/* Find subnet config file option */
	for (int i = 0; i < opt_count; i++) {
		if (strcmp(opts[i].key, "subnet_acl_file") == 0) {
			config_file = opts[i].value;
			break;
		}
	}

	pdata = calloc(1, sizeof(plugin_data_t));
	if (pdata == NULL)
		return MOSQ_ERR_NOMEM;

	pdata->identifier = identifier;

	/* Store config file path for reload */
	if (config_file != NULL) {
		pdata->config_file = strdup(config_file);
		if (pdata->config_file == NULL) {
			free(pdata);
			return MOSQ_ERR_NOMEM;
		}
	} else {
		pdata->config_file = NULL;
	}

	/* Load subnet ACL configuration */
	if (load_subnet_acl_config(pdata, config_file) != 0) {
		free(pdata->config_file);
		free(pdata);
		return MOSQ_ERR_UNKNOWN;
	}

	/* Register authentication callback only - subnet check is done during auth */
	rc = mosquitto_callback_register(identifier, MOSQ_EVT_BASIC_AUTH,
	                                 basic_auth_callback, NULL, pdata);
	if (rc != MOSQ_ERR_SUCCESS) {
		mosquitto_log_printf(MOSQ_LOG_ERR,
			"subnet_acl: Failed to register authentication callback");
		free(pdata->config_file);
		free(pdata);
		return rc;
	}

	/* Register reload callback to handle SIGHUP */
	rc = mosquitto_callback_register(identifier, MOSQ_EVT_RELOAD,
	                                 reload_callback, NULL, pdata);
	if (rc != MOSQ_ERR_SUCCESS) {
		mosquitto_log_printf(MOSQ_LOG_ERR,
			"subnet_acl: Failed to register reload callback");
		mosquitto_callback_unregister(identifier, MOSQ_EVT_BASIC_AUTH,
		                              basic_auth_callback, NULL);
		free(pdata->config_file);
		free(pdata);
		return rc;
	}

	mosquitto_log_printf(MOSQ_LOG_INFO,
		"subnet_acl: Plugin initialized with %d user(s)", pdata->user_count);

	/* Only assign user_data after all possible error paths */
	*user_data = pdata;

	return MOSQ_ERR_SUCCESS;
}

int mosquitto_plugin_cleanup(void *user_data,
			     struct mosquitto_opt *opts, int opt_count)
{
	plugin_data_t *pdata = user_data;

	if (pdata) {
		mosquitto_callback_unregister(pdata->identifier, MOSQ_EVT_BASIC_AUTH,
		                              basic_auth_callback, NULL);
		mosquitto_callback_unregister(pdata->identifier, MOSQ_EVT_RELOAD,
		                              reload_callback, NULL);
		free(pdata->config_file);
		free(pdata);
	}

	return MOSQ_ERR_SUCCESS;
}
