/*
 *
 *  Wireless daemon for Linux
 *
 *  Copyright (C) 2015  Intel Corporation. All rights reserved.
 *
 *  This library is free software; you can redistribute it and/or
 *  modify it under the terms of the GNU Lesser General Public
 *  License as published by the Free Software Foundation; either
 *  version 2.1 of the License, or (at your option) any later version.
 *
 *  This library is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 *  Lesser General Public License for more details.
 *
 *  You should have received a copy of the GNU Lesser General Public
 *  License along with this library; if not, write to the Free Software
 *  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 *
 */

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <stdbool.h>
#include <stdarg.h>
#include <errno.h>

#include <ell/ell.h>

#include "wscutil.h"

static const unsigned char wfa_ext[3] = { 0x00, 0x37, 0x2a };

void wsc_wfa_ext_iter_init(struct wsc_wfa_ext_iter *iter,
				const unsigned char *pdu, unsigned short len)
{
	iter->pdu = pdu;
	iter->max = len;
	iter->pos = 0;
}

bool wsc_wfa_ext_iter_next(struct wsc_wfa_ext_iter *iter)
{
	const unsigned char *start = iter->pdu + iter->pos;
	const unsigned char *end = iter->pdu + iter->max;
	unsigned char type;
	unsigned char len;

	if (iter->pos + 2 >= iter->max)
		return false;

	type = *start;
	start += 1;

	len = *start;
	start += 1;

	if (start + len > end)
		return false;

	iter->type = type;
	iter->len = len;
	iter->data = start;

	iter->pos = start + len - iter->pdu;

	return true;
}

void wsc_attr_iter_init(struct wsc_attr_iter *iter, const unsigned char *pdu,
			unsigned int len)
{
	iter->pdu = pdu;
	iter->max = len;
	iter->pos = 0;
}

bool wsc_attr_iter_next(struct wsc_attr_iter *iter)
{
	const unsigned char *start = iter->pdu + iter->pos;
	const unsigned char *end = iter->pdu + iter->max;
	unsigned short type;
	unsigned short len;

	/* Make sure we have at least type + len fields */
	if (iter->pos + 4 >= iter->max)
		return false;

	type = l_get_be16(start);
	start += 2;

	len = l_get_be16(start);
	start += 2;

	if (start + len > end)
		return false;

	iter->type = type;
	iter->len = len;
	iter->data = start;

	iter->pos = start + len - iter->pdu;

	return true;
}

bool wsc_attr_iter_recurse_wfa_ext(struct wsc_attr_iter *iter,
					struct wsc_wfa_ext_iter *wfa_iter)
{
	if (iter->type != WSC_ATTR_VENDOR_EXTENSION)
		return false;

	if (iter->len < 3)
		return false;

	if (memcmp(iter->data, wfa_ext, sizeof(wfa_ext)))
		return false;

	wsc_wfa_ext_iter_init(wfa_iter, iter->data + 3, iter->len - 3);

	return true;
}

enum attr_flag {
	ATTR_FLAG_REQUIRED  = 0x1,  /* Always required */
	ATTR_FLAG_VERSION2  = 0x2,  /* Included if Version2 is present */
	ATTR_FLAG_REGISTRAR = 0x4,  /* Included if Selected Registrar is true */
};

typedef bool (*attr_handler)(struct wsc_attr_iter *, void *);

static bool extract_uint8(struct wsc_attr_iter *iter, void *data)
{
	uint8_t *to = data;

	if (wsc_attr_iter_get_length(iter) != 1)
		return false;

	*to = *wsc_attr_iter_get_data(iter);

	return true;
}

static bool extract_uint16(struct wsc_attr_iter *iter, void *data)
{
	uint16_t *to = data;

	if (wsc_attr_iter_get_length(iter) != 2)
		return false;

	*to = l_get_be16(wsc_attr_iter_get_data(iter));

	return true;
}

static bool extract_bool(struct wsc_attr_iter *iter, void *data)
{
	bool *to = data;

	if (wsc_attr_iter_get_length(iter) != 1)
		return false;

	*to = *wsc_attr_iter_get_data(iter) ? true : false;

	return true;
}

static bool extract_uuid(struct wsc_attr_iter *iter, void *data)
{
	if (wsc_attr_iter_get_length(iter) != 16)
		return false;

	memcpy(data, wsc_attr_iter_get_data(iter), 16);
	return true;
}

static bool extract_ascii_string(struct wsc_attr_iter *iter, void *data,
					unsigned int max_len)
{
	char *out = data;
	const uint8_t *p;
	unsigned int len;
	unsigned int i;

	len = wsc_attr_iter_get_length(iter);
	if (len > max_len)
		return false;

	p = wsc_attr_iter_get_data(iter);

	for (i = 0; i < len; i++) {
		if (!p[i])
			break;

		if (!l_ascii_isprint(p[i]))
			return false;
	}

	memcpy(out, p, i);
	out[i] = '\0';
	return true;
}

static bool extract_utf8_string(struct wsc_attr_iter *iter, void *data,
					unsigned int max_len)
{
	char *out = data;
	const uint8_t *p;
	unsigned int len;
	unsigned int i;

	len = wsc_attr_iter_get_length(iter);
	if (len > max_len)
		return false;

	p = wsc_attr_iter_get_data(iter);

	for (i = 0; i < len; i++) {
		if (!p[i])
			break;
	}

	if (!l_utf8_validate((const char *) p, i, NULL))
		return false;

	memcpy(out, p, i);
	out[i] = '\0';
	return true;
}

static bool extract_nonce(struct wsc_attr_iter *iter, void *data)
{
	if (wsc_attr_iter_get_length(iter) != 16)
		return false;

	memcpy(data, wsc_attr_iter_get_data(iter), 16);
	return true;
}

static bool extract_association_state(struct wsc_attr_iter *iter, void *data)
{
	enum wsc_association_state *out = data;
	uint16_t as;

	if (!extract_uint16(iter, &as))
		return false;

	/* WSC 2.0.5: Table 31 */
	if (as > 4)
		return false;

	*out = as;
	return true;
}

static bool extract_authenticator(struct wsc_attr_iter *iter, void *data)
{
	if (wsc_attr_iter_get_length(iter) != 8)
		return false;

	memcpy(data, wsc_attr_iter_get_data(iter), 8);
	return true;
}

static bool extract_configuration_error(struct wsc_attr_iter *iter, void *data)
{
	enum wsc_configuration_error *out = data;
	uint16_t ce;

	if (!extract_uint16(iter, &ce))
		return false;

	/* WSC 2.0.5: Table 34 */
	if (ce > 20)
		return false;

	*out = ce;
	return true;
}

static bool extract_device_name(struct wsc_attr_iter *iter, void *data)
{
	return extract_utf8_string(iter, data, 32);
}

static bool extract_device_password_id(struct wsc_attr_iter *iter, void *data)
{
	uint16_t v;
	enum wsc_device_password_id *out = data;

	if (wsc_attr_iter_get_length(iter) != 2)
		return false;

	v = l_get_be16(wsc_attr_iter_get_data(iter));
	if (v > 0x0008)
		return false;

	*out = v;
	return true;
}

static bool extract_mac_address(struct wsc_attr_iter *iter, void *data)
{
	if (wsc_attr_iter_get_length(iter) != 6)
		return false;

	memcpy(data, wsc_attr_iter_get_data(iter), 6);
	return true;
}

static bool extract_manufacturer(struct wsc_attr_iter *iter, void *data)
{
	return extract_ascii_string(iter, data, 64);
}

static bool extract_message_type(struct wsc_attr_iter *iter, void *data)
{
	enum wsc_message_type *out = data;
	uint8_t mt;

	if (!extract_uint8(iter, &mt))
		return false;

	/* WSC 2.0.5: Table 42 */
	if (!mt || mt > 0x0f)
		return false;

	*out = mt;
	return true;
}

static bool extract_model_name(struct wsc_attr_iter *iter, void *data)
{
	return extract_ascii_string(iter, data, 32);
}

static bool extract_model_number(struct wsc_attr_iter *iter, void *data)
{
	return extract_ascii_string(iter, data, 32);
}

static bool extract_os_version(struct wsc_attr_iter *iter, void *data)
{
	uint32_t v;
	uint32_t *out = data;

	if (wsc_attr_iter_get_length(iter) != 4)
		return false;

	v = l_get_be32(wsc_attr_iter_get_data(iter));

	/*
	 * The OS Version component indicates what operating system is running
	 * on the device. It is a four-byte field. The most significant bit is
	 * reserved and always set to one.
	 */
	if ((v & 0x80000000) == 0)
		return false;

	*out = v & 0x7fffffff;
	return true;
}

static bool extract_public_key(struct wsc_attr_iter *iter, void *data)
{
	if (wsc_attr_iter_get_length(iter) != 192)
		return false;

	memcpy(data, wsc_attr_iter_get_data(iter), 192);
	return true;
}

static bool extract_primary_device_type(struct wsc_attr_iter *iter, void *data)
{
	struct wsc_primary_device_type *out = data;
	const uint8_t *p;
	uint16_t category;

	if (wsc_attr_iter_get_length(iter) != 8)
		return false;

	p = wsc_attr_iter_get_data(iter);
	category = l_get_be16(p);

	if (category > 12 && category != 255)
		return false;

	out->category = category;
	memcpy(out->oui, p + 2, 3);
	out->oui_type = p[5];
	out->subcategory = l_get_be16(p + 6);
	return true;
}

static bool extract_request_type(struct wsc_attr_iter *iter, void *data)
{
	enum wsc_request_type *out = data;
	uint8_t rt;

	if (!extract_uint8(iter, &rt))
		return false;

	/* WSC 2.0.5: Table 42 */
	if (rt > 3)
		return false;

	*out = rt;
	return true;
}

static bool extract_response_type(struct wsc_attr_iter *iter, void *data)
{
	enum wsc_response_type *out = data;
	uint8_t rt;

	if (!extract_uint8(iter, &rt))
		return false;

	/* WSC 2.0.5: Table 43 */
	if (rt > 3)
		return false;

	*out = rt;
	return true;
}

static bool extract_serial_number(struct wsc_attr_iter *iter, void *data)
{
	return extract_ascii_string(iter, data, 32);
}

static bool extract_version(struct wsc_attr_iter *iter, void *data)
{
	uint8_t *out = data;
	uint8_t v;

	if (!extract_uint8(iter, &v))
		return false;

	/*
	 * "This attribute is always set to value 0x10 (version 1.0)
	 * for backwards compatibility"
	 */
	if (v != 0x10)
		return false;

	*out = v;
	return true;
}

static bool extract_wsc_state(struct wsc_attr_iter *iter, void *data)
{
	uint8_t *out = data;
	uint8_t st;

	if (!extract_uint8(iter, &st))
		return false;

	if (st < 1 || st > 2)
		return false;

	*out = st;
	return true;
}

static attr_handler handler_for_type(enum wsc_attr type)
{
	switch (type) {
	case WSC_ATTR_AP_SETUP_LOCKED:
		return extract_bool;
	case WSC_ATTR_ASSOCIATION_STATE:
		return extract_association_state;
	case WSC_ATTR_AUTHENTICATION_TYPE_FLAGS:
		return extract_uint16;
	case WSC_ATTR_AUTHENTICATOR:
		return extract_authenticator;
	case WSC_ATTR_CONFIGURATION_ERROR:
		return extract_configuration_error;
	case WSC_ATTR_CONFIGURATION_METHODS:
		return extract_uint16;
	case WSC_ATTR_CONNECTION_TYPE_FLAGS:
		return extract_uint8;
	case WSC_ATTR_DEVICE_NAME:
		return extract_device_name;
	case WSC_ATTR_DEVICE_PASSWORD_ID:
		return extract_device_password_id;
	case WSC_ATTR_ENCRYPTION_TYPE_FLAGS:
		return extract_uint16;
	case WSC_ATTR_ENROLLEE_NONCE:
		return extract_nonce;
	case WSC_ATTR_MAC_ADDRESS:
		return extract_mac_address;
	case WSC_ATTR_MANUFACTURER:
		return extract_manufacturer;
	case WSC_ATTR_MESSAGE_TYPE:
		return extract_message_type;
	case WSC_ATTR_MODEL_NAME:
		return extract_model_name;
	case WSC_ATTR_MODEL_NUMBER:
		return extract_model_number;
	case WSC_ATTR_OS_VERSION:
		return extract_os_version;
	case WSC_ATTR_PUBLIC_KEY:
		return extract_public_key;
	case WSC_ATTR_PRIMARY_DEVICE_TYPE:
		return extract_primary_device_type;
	case WSC_ATTR_RF_BANDS:
		return extract_uint8;
	case WSC_ATTR_REGISTRAR_NONCE:
		return extract_nonce;
	case WSC_ATTR_REQUEST_TYPE:
		return extract_request_type;
	case WSC_ATTR_REQUESTED_DEVICE_TYPE:
		return extract_primary_device_type;
	case WSC_ATTR_RESPONSE_TYPE:
		return extract_response_type;
	case WSC_ATTR_SELECTED_REGISTRAR:
		return extract_bool;
	case WSC_ATTR_SELECTED_REGISTRAR_CONFIGURATION_METHODS:
		return extract_uint16;
	case WSC_ATTR_SERIAL_NUMBER:
		return extract_serial_number;
	case WSC_ATTR_VERSION:
		return extract_version;
	case WSC_ATTR_WSC_STATE:
		return extract_wsc_state;
	case WSC_ATTR_UUID_E:
		return extract_uuid;
	case WSC_ATTR_UUID_R:
		return extract_uuid;
	default:
		break;
	}

	return NULL;
}

struct attr_handler_entry {
	enum wsc_attr type;
	unsigned int flags;
	void *data;
	bool present;
};

static bool verify_version2(struct wsc_wfa_ext_iter *ext_iter)
{
	if (!wsc_wfa_ext_iter_next(ext_iter))
		return false;

	if (wsc_wfa_ext_iter_get_type(ext_iter) != WSC_WFA_EXTENSION_VERSION2)
		return false;

	if (wsc_wfa_ext_iter_get_length(ext_iter) != 1)
		return false;

	return true;
}

static int wsc_parse_attrs(const unsigned char *pdu, unsigned int len,
				bool *out_version2,
				struct wsc_wfa_ext_iter *ext_iter,
				uint8_t *authenticator,
				enum wsc_attr type, ...)
{
	struct wsc_attr_iter iter;
	struct l_queue *entries;
	const struct l_queue_entry *e;
	va_list args;
	bool version2 = false;
	bool sr = false;
	bool have_required = true;
	bool parse_error = false;

	wsc_attr_iter_init(&iter, pdu, len);

	va_start(args, type);

	entries = l_queue_new();

	while (type != WSC_ATTR_INVALID) {
		struct attr_handler_entry *entry;

		entry = l_new(struct attr_handler_entry, 1);

		entry->type = type;
		entry->flags = va_arg(args, unsigned int);
		entry->data = va_arg(args, void *);

		type = va_arg(args, enum wsc_attr);
		l_queue_push_tail(entries, entry);
	}

	va_end(args);
	e = l_queue_get_entries(entries);

	while (wsc_attr_iter_next(&iter)) {
		attr_handler handler;
		struct attr_handler_entry *entry;
		const struct l_queue_entry *e2;

		for (e2 = e; e2; e2 = e2->next) {
			entry = e2->data;

			if (wsc_attr_iter_get_type(&iter) == entry->type) {
				entry->present = true;
				break;
			}

			if (entry->flags & ATTR_FLAG_REQUIRED) {
				have_required = false;
				goto done;
			}
		}

		if (e2 == NULL) {
			if (wsc_attr_iter_get_type(&iter)
					!= WSC_ATTR_VENDOR_EXTENSION)
				break;

			if (!wsc_attr_iter_recurse_wfa_ext(&iter, ext_iter))
				break;

			if (!verify_version2(ext_iter)) {
				parse_error = true;
				goto done;
			}

			version2 = true;
			continue;
		}

		if (entry->type == WSC_ATTR_SELECTED_REGISTRAR)
			sr = true;

		handler = handler_for_type(entry->type);

		if (!handler(&iter, entry->data)) {
			parse_error = true;
			goto done;
		}

		e = e2->next;
	}

	for (; e; e = e->next) {
		struct attr_handler_entry *entry = e->data;

		if (entry->flags & ATTR_FLAG_REQUIRED) {
			parse_error = true;
			goto done;
		}
	}

	/* Authenticator element must be the last element */
	if (authenticator) {
		while (wsc_attr_iter_get_type(&iter) !=
						WSC_ATTR_AUTHENTICATOR) {
			if (!wsc_attr_iter_next(&iter)) {
				have_required = false;
				goto done;
			}
		}

		if (!extract_authenticator(&iter, authenticator)) {
			parse_error = true;
			goto done;
		}

		if (wsc_attr_iter_next(&iter) != false) {
			parse_error = true;
			goto done;
		}
	}

	/*
	 * Check for Version 2.0 required attributes.
	 * If version2 attribute is present in the WFA Vendor field,
	 * then check the required attributes are present.  Mostly relevant
	 * for Probe Request messages according to 8.2.4 in WSC 2.0.5
	 */
	if (version2) {
		struct attr_handler_entry *entry;

		for (e = l_queue_get_entries(entries); e; e = e->next) {
			entry = e->data;

			if (!(entry->flags & ATTR_FLAG_VERSION2))
				continue;

			if (entry->present)
				continue;

			parse_error = true;
			goto done;
		}
	}

	/*
	 * If Selected Registrar is present and true, then certain attributes
	 * must also be present.
	 */
	if (sr) {
		struct attr_handler_entry *entry;

		for (e = l_queue_get_entries(entries); e; e = e->next) {
			entry = e->data;

			if (!(entry->flags & ATTR_FLAG_REGISTRAR))
				continue;

			if (entry->present)
				continue;

			parse_error = true;
			goto done;
		}

	}

done:
	l_queue_destroy(entries, l_free);

	if (!have_required)
		return -EINVAL;
	if (parse_error)
		return -EBADMSG;

	*out_version2 = version2;

	return 0;
}

static bool wfa_extract_bool(struct wsc_wfa_ext_iter *iter, void *data)
{
	bool *to = data;

	if (wsc_wfa_ext_iter_get_length(iter) != 1)
		return false;

	*to = *wsc_wfa_ext_iter_get_data(iter);
	return true;
}

static bool wfa_extract_authorized_macs(struct wsc_wfa_ext_iter *iter,
								void *data)
{
	uint8_t *to = data;
	unsigned int len = wsc_wfa_ext_iter_get_length(iter);
	unsigned int mod;

	if (!len || len > 30)
		return false;

	mod = len % 6;
	if (mod)
		return false;

	memcpy(to, wsc_wfa_ext_iter_get_data(iter), len);
	return true;
}

static bool wfa_extract_registrar_configuration_methods(
				struct wsc_wfa_ext_iter *iter, void *data)
{
	uint16_t *to = data;

	if (wsc_wfa_ext_iter_get_length(iter) != 2)
		return false;

	*to = l_get_be16(wsc_wfa_ext_iter_get_data(iter));
	return true;
}

#define REQUIRED(attr, out) \
	WSC_ATTR_ ## attr, ATTR_FLAG_REQUIRED, out

#define OPTIONAL(attr, out) \
	WSC_ATTR_ ## attr, 0, out

#define REGISTRAR(attr, out) \
	WSC_ATTR_ ## attr, ATTR_FLAG_REGISTRAR, out

#define VERSION2(attr, out) \
	WSC_ATTR_ ## attr, ATTR_FLAG_VERSION2, out

int wsc_parse_beacon(const unsigned char *pdu, unsigned int len,
				struct wsc_beacon *out)
{
	int r;
	struct wsc_wfa_ext_iter iter;
	uint8_t version;

	memset(out, 0, sizeof(struct wsc_beacon));

	r = wsc_parse_attrs(pdu, len, &out->version2, &iter, NULL,
		REQUIRED(VERSION, &version),
		REQUIRED(WSC_STATE, &out->state),
		OPTIONAL(AP_SETUP_LOCKED, &out->ap_setup_locked),
		OPTIONAL(SELECTED_REGISTRAR, &out->selected_registrar),
		REGISTRAR(DEVICE_PASSWORD_ID, &out->device_password_id),
		REGISTRAR(SELECTED_REGISTRAR_CONFIGURATION_METHODS,
					&out->selected_reg_config_methods),
		OPTIONAL(UUID_E, &out->uuid_e),
		OPTIONAL(RF_BANDS, &out->rf_bands),
		WSC_ATTR_INVALID);

	if (r < 0)
		return r;

	if (!wsc_wfa_ext_iter_next(&iter))
		goto done;

	if (wsc_wfa_ext_iter_get_type(&iter) ==
					WSC_WFA_EXTENSION_AUTHORIZED_MACS) {
		if (!wfa_extract_authorized_macs(&iter, &out->authorized_macs))
			return -EBADMSG;

		if (!wsc_wfa_ext_iter_next(&iter))
			goto done;
	}

	if (wsc_wfa_ext_iter_get_type(&iter) ==
			WSC_WFA_EXTENSION_REGISTRAR_CONFIGRATION_METHODS) {
		if (!wfa_extract_registrar_configuration_methods(&iter,
						&out->reg_config_methods))
			return -EBADMSG;

		if (!wsc_wfa_ext_iter_next(&iter))
			goto done;
	}

	return -EINVAL;

done:
	return 0;
}

int wsc_parse_probe_response(const unsigned char *pdu, unsigned int len,
				struct wsc_probe_response *out)
{
	int r;
	struct wsc_wfa_ext_iter iter;
	uint8_t version;

	memset(out, 0, sizeof(struct wsc_probe_response));

	r = wsc_parse_attrs(pdu, len, &out->version2, &iter, NULL,
		REQUIRED(VERSION, &version),
		REQUIRED(WSC_STATE, &out->state),
		OPTIONAL(AP_SETUP_LOCKED, &out->ap_setup_locked),
		OPTIONAL(SELECTED_REGISTRAR, &out->selected_registrar),
		REGISTRAR(DEVICE_PASSWORD_ID, &out->device_password_id),
		REGISTRAR(SELECTED_REGISTRAR_CONFIGURATION_METHODS,
					&out->selected_reg_config_methods),
		REQUIRED(RESPONSE_TYPE, &out->response_type),
		REQUIRED(UUID_E, &out->uuid_e),
		REQUIRED(MANUFACTURER, &out->manufacturer),
		REQUIRED(MODEL_NAME, &out->model_name),
		REQUIRED(MODEL_NUMBER, &out->model_number),
		REQUIRED(SERIAL_NUMBER, &out->serial_number),
		REQUIRED(PRIMARY_DEVICE_TYPE, &out->primary_device_type),
		REQUIRED(DEVICE_NAME, &out->device_name),
		REQUIRED(CONFIGURATION_METHODS, &out->config_methods),
		OPTIONAL(RF_BANDS, &out->rf_bands),
		WSC_ATTR_INVALID);

	if (r < 0)
		return r;

	if (!wsc_wfa_ext_iter_next(&iter))
		goto done;

	if (wsc_wfa_ext_iter_get_type(&iter) ==
					WSC_WFA_EXTENSION_AUTHORIZED_MACS) {
		if (!wfa_extract_authorized_macs(&iter, &out->authorized_macs))
			return -EBADMSG;

		if (!wsc_wfa_ext_iter_next(&iter))
			goto done;
	}

	if (wsc_wfa_ext_iter_get_type(&iter) ==
			WSC_WFA_EXTENSION_REGISTRAR_CONFIGRATION_METHODS) {
		if (!wfa_extract_registrar_configuration_methods(&iter,
						&out->reg_config_methods))
			return -EBADMSG;

		if (!wsc_wfa_ext_iter_next(&iter))
			goto done;
	}

	return -EINVAL;

done:
	return 0;
}

int wsc_parse_probe_request(const unsigned char *pdu, unsigned int len,
				struct wsc_probe_request *out)
{
	int r;
	struct wsc_wfa_ext_iter iter;
	uint8_t version;

	memset(out, 0, sizeof(struct wsc_probe_request));

	r = wsc_parse_attrs(pdu, len, &out->version2, &iter, NULL,
		REQUIRED(VERSION, &version),
		REQUIRED(REQUEST_TYPE, &out->request_type),
		REQUIRED(CONFIGURATION_METHODS, &out->config_methods),
		REQUIRED(UUID_E, &out->uuid_e),
		REQUIRED(PRIMARY_DEVICE_TYPE, &out->primary_device_type),
		REQUIRED(RF_BANDS, &out->rf_bands),
		REQUIRED(ASSOCIATION_STATE, &out->association_state),
		REQUIRED(CONFIGURATION_ERROR, &out->configuration_error),
		REQUIRED(DEVICE_PASSWORD_ID, &out->device_password_id),
		VERSION2(MANUFACTURER, &out->manufacturer),
		VERSION2(MODEL_NAME, &out->model_name),
		VERSION2(MODEL_NUMBER, &out->model_number),
		VERSION2(DEVICE_NAME, &out->device_name),
		OPTIONAL(REQUESTED_DEVICE_TYPE, &out->requested_device_type),
		WSC_ATTR_INVALID);

	if (r < 0)
		return r;

	if (!wsc_wfa_ext_iter_next(&iter))
		goto done;

	if (wsc_wfa_ext_iter_get_type(&iter) ==
					WSC_WFA_EXTENSION_REQUEST_TO_ENROLL) {
		if (!wfa_extract_bool(&iter, &out->request_to_enroll))
			return -EBADMSG;

		if (!wsc_wfa_ext_iter_next(&iter))
			goto done;
	}

	return -EINVAL;

done:
	return 0;
}

int wsc_parse_m1(const uint8_t *pdu, uint32_t len, struct wsc_m1 *out)
{
	int r;
	struct wsc_wfa_ext_iter iter;
	uint8_t version;
	enum wsc_message_type msg_type;

	memset(out, 0, sizeof(struct wsc_m1));

	r = wsc_parse_attrs(pdu, len, &out->version2, &iter, NULL,
		REQUIRED(VERSION, &version),
		REQUIRED(MESSAGE_TYPE, &msg_type),
		REQUIRED(UUID_E, &out->uuid_e),
		REQUIRED(MAC_ADDRESS, &out->addr),
		REQUIRED(ENROLLEE_NONCE, &out->enrollee_nonce),
		REQUIRED(PUBLIC_KEY, &out->public_key),
		REQUIRED(AUTHENTICATION_TYPE_FLAGS, &out->auth_type_flags),
		REQUIRED(ENCRYPTION_TYPE_FLAGS, &out->encryption_type_flags),
		REQUIRED(CONNECTION_TYPE_FLAGS, &out->connection_type_flags),
		REQUIRED(CONFIGURATION_METHODS, &out->config_methods),
		REQUIRED(WSC_STATE, &out->state),
		REQUIRED(MANUFACTURER, &out->manufacturer),
		REQUIRED(MODEL_NAME, &out->model_name),
		REQUIRED(MODEL_NUMBER, &out->model_number),
		REQUIRED(SERIAL_NUMBER, &out->serial_number),
		REQUIRED(PRIMARY_DEVICE_TYPE, &out->primary_device_type),
		REQUIRED(DEVICE_NAME, &out->device_name),
		REQUIRED(RF_BANDS, &out->rf_bands),
		REQUIRED(ASSOCIATION_STATE, &out->association_state),
		REQUIRED(DEVICE_PASSWORD_ID, &out->device_password_id),
		REQUIRED(CONFIGURATION_ERROR, &out->configuration_error),
		REQUIRED(OS_VERSION, &out->os_version),
		WSC_ATTR_INVALID);

	if (r < 0)
		return r;

	if (msg_type != WSC_MESSAGE_TYPE_M1)
		return -EBADMSG;

	/* WSC 2.0.5, Section 8.3.1: "Specific RF band used for this message" */
	if (__builtin_popcount(out->rf_bands) != 1)
		return -EBADMSG;

	if (!wsc_wfa_ext_iter_next(&iter))
		goto done;

	if (wsc_wfa_ext_iter_get_type(&iter) ==
					WSC_WFA_EXTENSION_REQUEST_TO_ENROLL) {
		if (!wfa_extract_bool(&iter, &out->request_to_enroll))
			return -EBADMSG;

		if (!wsc_wfa_ext_iter_next(&iter))
			goto done;
	}

	return -EINVAL;

done:
	return 0;
}

int wsc_parse_m2(const uint8_t *pdu, uint32_t len, struct wsc_m2 *out)
{
	int r;
	struct wsc_wfa_ext_iter iter;
	uint8_t version;
	enum wsc_message_type msg_type;

	memset(out, 0, sizeof(struct wsc_m2));

	r = wsc_parse_attrs(pdu, len, &out->version2, &iter, out->authenticator,
		REQUIRED(VERSION, &version),
		REQUIRED(MESSAGE_TYPE, &msg_type),
		REQUIRED(ENROLLEE_NONCE, &out->enrollee_nonce),
		REQUIRED(REGISTRAR_NONCE, &out->registrar_nonce),
		REQUIRED(UUID_R, &out->uuid_r),
		REQUIRED(PUBLIC_KEY, &out->public_key),
		REQUIRED(AUTHENTICATION_TYPE_FLAGS, &out->auth_type_flags),
		REQUIRED(ENCRYPTION_TYPE_FLAGS, &out->encryption_type_flags),
		REQUIRED(CONNECTION_TYPE_FLAGS, &out->connection_type_flags),
		REQUIRED(CONFIGURATION_METHODS, &out->config_methods),
		REQUIRED(MANUFACTURER, &out->manufacturer),
		REQUIRED(MODEL_NAME, &out->model_name),
		REQUIRED(MODEL_NUMBER, &out->model_number),
		REQUIRED(SERIAL_NUMBER, &out->serial_number),
		REQUIRED(PRIMARY_DEVICE_TYPE, &out->primary_device_type),
		REQUIRED(DEVICE_NAME, &out->device_name),
		REQUIRED(RF_BANDS, &out->rf_bands),
		REQUIRED(ASSOCIATION_STATE, &out->association_state),
		REQUIRED(CONFIGURATION_ERROR, &out->configuration_error),
		REQUIRED(DEVICE_PASSWORD_ID, &out->device_password_id),
		REQUIRED(OS_VERSION, &out->os_version),
		WSC_ATTR_INVALID);

	if (r < 0)
		return r;

	if (msg_type != WSC_MESSAGE_TYPE_M2)
		return -EBADMSG;

	/* WSC 2.0.5, Section 8.3.2: "Specific RF band used for this message" */
	if (__builtin_popcount(out->rf_bands) != 1)
		return -EBADMSG;

	return 0;
}

struct wsc_attr_builder {
	size_t capacity;
	uint8_t *buf;
	size_t offset;
	uint16_t curlen;
};

static void wsc_attr_builder_grow(struct wsc_attr_builder *builder)
{
	builder->buf = l_realloc(builder->buf, builder->capacity * 2);
	builder->capacity *= 2;
}

static bool wsc_attr_builder_start_attr(struct wsc_attr_builder *builder,
						enum wsc_attr type)
{
	uint8_t *bytes;

	/* TLVs must be length > 0 */
	if (builder->curlen == 0 && builder->offset != 0)
		return false;

	/* Record previous attribute's length */
	if (builder->curlen > 0) {
		bytes = builder->buf + builder->offset;
		l_put_be16(builder->curlen, bytes + 2);
		builder->offset += 4 + builder->curlen;
		builder->curlen = 0;
	}

	if (builder->offset + 4 >= builder->capacity)
		wsc_attr_builder_grow(builder);

	bytes = builder->buf + builder->offset;
	l_put_be16(type, bytes);

	return true;
}

static bool wsc_attr_builder_put_u8(struct wsc_attr_builder *builder, uint8_t v)
{
	if (builder->offset + 4 + builder->curlen + 1 >= builder->capacity)
		wsc_attr_builder_grow(builder);

	builder->buf[builder->offset + 4 + builder->curlen] = v;
	builder->curlen += 1;

	return true;
}

static bool wsc_attr_builder_put_u16(struct wsc_attr_builder *builder,
								uint16_t v)
{
	if (builder->offset + 4 + builder->curlen + 2 >= builder->capacity)
		wsc_attr_builder_grow(builder);

	l_put_be16(v, builder->buf + builder->offset + 4 + builder->curlen);
	builder->curlen += 2;

	return true;
}

static bool wsc_attr_builder_put_u32(struct wsc_attr_builder *builder,
								uint32_t v)
{
	if (builder->offset + 4 + builder->curlen + 4 >= builder->capacity)
		wsc_attr_builder_grow(builder);

	l_put_be32(v, builder->buf + builder->offset + 4 + builder->curlen);
	builder->curlen += 4;

	return true;
}

static bool wsc_attr_builder_put_bytes(struct wsc_attr_builder *builder,
					const void *bytes, size_t size)
{
	while (builder->offset + 4 + builder->curlen + size >=
							builder->capacity)
		wsc_attr_builder_grow(builder);

	memcpy(builder->buf + builder->offset + 4 + builder->curlen,
								bytes, size);
	builder->curlen += size;

	return true;
}

static bool wsc_attr_builder_put_oui(struct wsc_attr_builder *builder,
							const uint8_t *oui)
{
	if (builder->offset + 4 + builder->curlen + 3 >= builder->capacity)
		wsc_attr_builder_grow(builder);

	memcpy(builder->buf + builder->offset + 4 + builder->curlen, oui, 3);
	builder->curlen += 3;

	return true;
}

static bool wsc_attr_builder_put_string(struct wsc_attr_builder *builder,
							const char *string)
{
	size_t len;

	len = string ? strlen(string) : 0;

	if (len == 0) {
		string = " ";
		len = 1;
	}

	if (builder->offset + 4 + builder->curlen + len >= builder->capacity)
		wsc_attr_builder_grow(builder);

	memcpy(builder->buf + builder->offset + 4 + builder->curlen,
								string, len);
	builder->curlen += len;

	return true;
}

static struct wsc_attr_builder *wsc_attr_builder_new(size_t initial_capacity)
{
	struct wsc_attr_builder *builder;

	if (initial_capacity == 0)
		return NULL;

	builder = l_new(struct wsc_attr_builder, 1);
	builder->buf = l_malloc(initial_capacity);
	builder->capacity = initial_capacity;

	return builder;
}

static uint8_t *wsc_attr_builder_free(struct wsc_attr_builder *builder,
					bool free_contents,
					size_t *out_size)
{
	uint8_t *ret;

	if (builder->curlen > 0) {
		uint8_t *bytes = builder->buf + builder->offset;
		l_put_be16(builder->curlen, bytes + 2);
		builder->offset += 4 + builder->curlen;
		builder->curlen = 0;
	}

	if (free_contents) {
		l_free(builder->buf);
		builder->buf = NULL;
	}

	ret = builder->buf;

	if (out_size)
		*out_size = builder->offset;

	l_free(builder);

	return ret;
}

static void build_association_state(struct wsc_attr_builder *builder,
					enum wsc_association_state state)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_ASSOCIATION_STATE);
	wsc_attr_builder_put_u16(builder, state);
}

static void build_authentication_type_flags(struct wsc_attr_builder *builder,
					uint16_t auth_type_flags)
{
	wsc_attr_builder_start_attr(builder,
					WSC_ATTR_AUTHENTICATION_TYPE_FLAGS);
	wsc_attr_builder_put_u16(builder, auth_type_flags);
}

static void build_authenticator(struct wsc_attr_builder *builder,
					const uint8_t *authenticator)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_AUTHENTICATOR);
	wsc_attr_builder_put_bytes(builder, authenticator, 8);
}

static void build_configuration_error(struct wsc_attr_builder *builder,
					enum wsc_configuration_error error)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_CONFIGURATION_ERROR);
	wsc_attr_builder_put_u16(builder, error);
}

static void build_configuration_methods(struct wsc_attr_builder *builder,
						uint16_t config_methods)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_CONFIGURATION_METHODS);
	wsc_attr_builder_put_u16(builder, config_methods);
}

static void build_connection_type_flags(struct wsc_attr_builder *builder,
						uint8_t connection_type_flags)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_CONNECTION_TYPE_FLAGS);
	wsc_attr_builder_put_u8(builder, connection_type_flags);
}

static void build_device_name(struct wsc_attr_builder *builder,
						const char *device_name)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_DEVICE_NAME);
	wsc_attr_builder_put_string(builder, device_name);
}

static void build_device_password_id(struct wsc_attr_builder *builder,
					enum wsc_device_password_id id)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_DEVICE_PASSWORD_ID);
	wsc_attr_builder_put_u16(builder, id);
}

static void build_encryption_type_flags(struct wsc_attr_builder *builder,
						uint16_t encryption_type_flags)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_ENCRYPTION_TYPE_FLAGS);
	wsc_attr_builder_put_u16(builder, encryption_type_flags);
}

static void build_enrollee_nonce(struct wsc_attr_builder *builder,
							const uint8_t *nonce)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_ENROLLEE_NONCE);
	wsc_attr_builder_put_bytes(builder, nonce, 16);
}

static void build_mac_address(struct wsc_attr_builder *builder,
							const uint8_t *addr)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_MAC_ADDRESS);
	wsc_attr_builder_put_bytes(builder, addr, 6);
}

static void build_manufacturer(struct wsc_attr_builder *builder,
						const char *manufacturer)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_MANUFACTURER);
	wsc_attr_builder_put_string(builder, manufacturer);
}

static void build_message_type(struct wsc_attr_builder *builder,
						enum wsc_message_type type)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_MESSAGE_TYPE);
	wsc_attr_builder_put_u8(builder, type);
}

static void build_model_name(struct wsc_attr_builder *builder,
						const char *model_name)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_MODEL_NAME);
	wsc_attr_builder_put_string(builder, model_name);
}

static void build_model_number(struct wsc_attr_builder *builder,
						const char *model_number)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_MODEL_NUMBER);
	wsc_attr_builder_put_string(builder, model_number);
}

static void build_os_version(struct wsc_attr_builder *builder,
							uint32_t os_version)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_OS_VERSION);
	wsc_attr_builder_put_u32(builder, os_version | 0x80000000);
}

static void build_primary_device_type(struct wsc_attr_builder *builder,
				const struct wsc_primary_device_type *pdt)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_PRIMARY_DEVICE_TYPE);
	wsc_attr_builder_put_u16(builder, pdt->category);
	wsc_attr_builder_put_oui(builder, pdt->oui);
	wsc_attr_builder_put_u8(builder, pdt->oui_type);
	wsc_attr_builder_put_u16(builder, pdt->subcategory);
}

static void build_public_key(struct wsc_attr_builder *builder,
						const uint8_t *public_key)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_PUBLIC_KEY);
	wsc_attr_builder_put_bytes(builder, public_key, 192);
}

static void build_registrar_nonce(struct wsc_attr_builder *builder,
							const uint8_t *nonce)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_REGISTRAR_NONCE);
	wsc_attr_builder_put_bytes(builder, nonce, 16);
}

static void build_request_type(struct wsc_attr_builder *builder,
						enum wsc_request_type type)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_REQUEST_TYPE);
	wsc_attr_builder_put_u8(builder, type);
}

static void build_rf_bands(struct wsc_attr_builder *builder, uint8_t rf_bands)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_RF_BANDS);
	wsc_attr_builder_put_u8(builder, rf_bands);
}

static void build_serial_number(struct wsc_attr_builder *builder,
						const char *serial_number)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_SERIAL_NUMBER);
	wsc_attr_builder_put_string(builder, serial_number);
}

static void build_uuid_e(struct wsc_attr_builder *builder, const uint8_t *uuid)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_UUID_E);
	wsc_attr_builder_put_bytes(builder, uuid, 16);
}

static void build_uuid_r(struct wsc_attr_builder *builder, const uint8_t *uuid)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_UUID_R);
	wsc_attr_builder_put_bytes(builder, uuid, 16);
}

static void build_version(struct wsc_attr_builder *builder, uint8_t version)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_VERSION);
	wsc_attr_builder_put_u8(builder, version);
}

static void build_wsc_state(struct wsc_attr_builder *builder,
					enum wsc_state state)
{
	wsc_attr_builder_start_attr(builder, WSC_ATTR_WSC_STATE);
	wsc_attr_builder_put_u8(builder, state);
}

uint8_t *wsc_build_probe_request(const struct wsc_probe_request *probe_request,
							size_t *out_len)
{
	struct wsc_attr_builder *builder;
	uint8_t *ret;

	builder = wsc_attr_builder_new(512);
	build_version(builder, 0x10);
	build_request_type(builder, probe_request->request_type);
	build_configuration_methods(builder, probe_request->config_methods);
	build_uuid_e(builder, probe_request->uuid_e);
	build_primary_device_type(builder, &probe_request->primary_device_type);
	build_rf_bands(builder, probe_request->rf_bands);
	build_association_state(builder, probe_request->association_state);
	build_configuration_error(builder, probe_request->configuration_error);
	build_device_password_id(builder, probe_request->device_password_id);

	if (!probe_request->version2)
		goto done;

	build_manufacturer(builder, probe_request->manufacturer);
	build_model_name(builder, probe_request->model_name);
	build_model_number(builder, probe_request->model_number);
	build_device_name(builder, probe_request->device_name);

	/* Put in the WFA Vendor Extension */
	wsc_attr_builder_start_attr(builder, WSC_ATTR_VENDOR_EXTENSION);
	wsc_attr_builder_put_oui(builder, wfa_ext);
	wsc_attr_builder_put_u8(builder, WSC_WFA_EXTENSION_VERSION2);
	wsc_attr_builder_put_u8(builder, 1);
	wsc_attr_builder_put_u8(builder, 0x20);

	if (!probe_request->request_to_enroll)
		goto done;

	wsc_attr_builder_put_u8(builder, WSC_WFA_EXTENSION_REQUEST_TO_ENROLL);
	wsc_attr_builder_put_u8(builder, 1);
	wsc_attr_builder_put_u8(builder, 1);

done:
	ret = wsc_attr_builder_free(builder, false, out_len);
	return ret;
}

uint8_t *wsc_build_m1(const struct wsc_m1 *m1, size_t *out_len)
{
	struct wsc_attr_builder *builder;
	uint8_t *ret;

	builder = wsc_attr_builder_new(1024);
	build_version(builder, 0x10);
	build_message_type(builder, WSC_MESSAGE_TYPE_M1);
	build_uuid_e(builder, m1->uuid_e);
	build_mac_address(builder, m1->addr);
	build_enrollee_nonce(builder, m1->enrollee_nonce);
	build_public_key(builder, m1->public_key);
	build_authentication_type_flags(builder, m1->auth_type_flags);
	build_encryption_type_flags(builder, m1->encryption_type_flags);
	build_connection_type_flags(builder, m1->connection_type_flags);
	build_configuration_methods(builder, m1->config_methods);
	build_wsc_state(builder, m1->state);
	build_manufacturer(builder, m1->manufacturer);
	build_model_name(builder, m1->model_name);
	build_model_number(builder, m1->model_number);
	build_serial_number(builder, m1->serial_number);
	build_primary_device_type(builder, &m1->primary_device_type);
	build_device_name(builder, m1->device_name);
	build_rf_bands(builder, m1->rf_bands);
	build_association_state(builder, m1->association_state);
	build_device_password_id(builder, m1->device_password_id);
	build_configuration_error(builder, m1->configuration_error);
	build_os_version(builder, m1->os_version);

	if (!m1->version2)
		goto done;

	/* Put in the WFA Vendor Extension */
	wsc_attr_builder_start_attr(builder, WSC_ATTR_VENDOR_EXTENSION);
	wsc_attr_builder_put_oui(builder, wfa_ext);
	wsc_attr_builder_put_u8(builder, WSC_WFA_EXTENSION_VERSION2);
	wsc_attr_builder_put_u8(builder, 1);
	wsc_attr_builder_put_u8(builder, 0x20);

	if (!m1->request_to_enroll)
		goto done;

	wsc_attr_builder_put_u8(builder, WSC_WFA_EXTENSION_REQUEST_TO_ENROLL);
	wsc_attr_builder_put_u8(builder, 1);
	wsc_attr_builder_put_u8(builder, 1);

done:
	ret = wsc_attr_builder_free(builder, false, out_len);
	return ret;
}

uint8_t *wsc_build_m2(const struct wsc_m2 *m2, size_t *out_len)
{
	struct wsc_attr_builder *builder;
	uint8_t *ret;

	builder = wsc_attr_builder_new(1024);
	build_version(builder, 0x10);
	build_message_type(builder, WSC_MESSAGE_TYPE_M2);
	build_enrollee_nonce(builder, m2->enrollee_nonce);
	build_registrar_nonce(builder, m2->registrar_nonce);
	build_uuid_r(builder, m2->uuid_r);
	build_public_key(builder, m2->public_key);
	build_authentication_type_flags(builder, m2->auth_type_flags);
	build_encryption_type_flags(builder, m2->encryption_type_flags);
	build_connection_type_flags(builder, m2->connection_type_flags);
	build_configuration_methods(builder, m2->config_methods);
	build_manufacturer(builder, m2->manufacturer);
	build_model_name(builder, m2->model_name);
	build_model_number(builder, m2->model_number);
	build_serial_number(builder, m2->serial_number);
	build_primary_device_type(builder, &m2->primary_device_type);
	build_device_name(builder, m2->device_name);
	build_rf_bands(builder, m2->rf_bands);
	build_association_state(builder, m2->association_state);
	build_configuration_error(builder, m2->configuration_error);
	build_device_password_id(builder, m2->device_password_id);
	build_os_version(builder, m2->os_version);

	if (!m2->version2)
		goto done;

	/* Put in the WFA Vendor Extension */
	wsc_attr_builder_start_attr(builder, WSC_ATTR_VENDOR_EXTENSION);
	wsc_attr_builder_put_oui(builder, wfa_ext);
	wsc_attr_builder_put_u8(builder, WSC_WFA_EXTENSION_VERSION2);
	wsc_attr_builder_put_u8(builder, 1);
	wsc_attr_builder_put_u8(builder, 0x20);

done:
	build_authenticator(builder, m2->authenticator);

	ret = wsc_attr_builder_free(builder, false, out_len);
	return ret;
}


bool wsc_uuid_from_addr(const uint8_t addr[], uint8_t *out_uuid)
{
	/* Reuse the NSID from WPA Supplicant for compatibility */
	static const uint8_t nsid[] = {
		0x52, 0x64, 0x80, 0xf8, 0xc9, 0x9b, 0x4b, 0xe5,
		0xa6, 0x55, 0x58, 0xed, 0x5f, 0x5d, 0x60, 0x84,
	};

	return l_uuid_v5(nsid, addr, 6, out_uuid);
}