/*
 * LibWPS
 * Copyright (c) 2011, Tactical Network Solutions, Craig Heffner <cheffner@tacnetsol.com>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2 as
 * published by the Free Software Foundation.
 *
 * See README and LICENSE for more details.
 */

#define LIBWPS_C

#include "libwps.h"

/*
 * This is the only function that external code should call. 
 *
 * const u_char *packet		Pointer to a beacon or probe response packet
 * size_t len			Size of the packet
 * struct libwps_data *wps	Pointer to an allocated libwps_data structure
 *
 * Returns 1 if WPS data was found and the libwps_data structure has been populated.
 * Returns 0 if no WPS data was found.
 */
int parse_wps_parameters(const u_char *packet, size_t len, struct libwps_data *wps)
{
	const u_char *data = NULL;
	size_t data_len = 0, offset = 0;
	struct radio_tap_header *rt_header = NULL;
	int ret_val = 0;

	if(wps)
	{
		memset(wps, 0, sizeof(struct libwps_data));
	
		if(len > (sizeof(struct radio_tap_header) + 
			  sizeof(struct dot11_frame_header) + 
		 	  sizeof(struct management_frame)))
		{
			rt_header = (struct radio_tap_header *) libwps_radio_header(packet, len);

			offset = rt_header->len + sizeof(struct dot11_frame_header) + sizeof(struct management_frame);
			data = (packet + offset);
			data_len = (len - offset);

			ret_val = parse_wps_tag(data, data_len, wps);
		}
	}

	return ret_val;
}

/* Parse and print WPS data in beacon packets and probe responses */
int parse_wps_tag(const u_char *tags, size_t len, struct libwps_data *wps)
{
	unsigned char *wps_ie_data = NULL, *el = NULL;
	char *ptr = NULL, *src = NULL;
	int i = 0, ret_val = 0;
	size_t wps_data_len = 0, el_len = 0;
	enum wps_el_number elements[] = {
			VERSION,
			STATE,
			LOCKED,
			MANUFACTURER,
			MODEL_NAME,
			MODEL_NUMBER,
			DEVICE_NAME,
			SSID,
			UUID,
			SERIAL,
			SELECTED_REGISTRAR,
			RESPONSE_TYPE,
			PRIMARY_DEVICE_TYPE,
			CONFIG_METHODS,
			RF_BANDS,
			OS_VERSION
	};

	/* Get the WPS IE data blob */
	wps_ie_data = get_wps_data(tags, len, &wps_data_len);
	wps->locked = UNSPECIFIED;

	if(wps_ie_data)
	{
		for(i=0; i<sizeof(elements)/sizeof(elements[0]); i++)
		{
			/* Search for each WPS element inside the WPS IE data blob */
			el = get_wps_data_element(wps_ie_data, wps_data_len, elements[i], &el_len);
			if(el)
			{
				if(el_len > LIBWPS_MAX_STR_LEN)
				{
					el_len = LIBWPS_MAX_STR_LEN;
				}

				switch(elements[i])
				{
					case VERSION:
						wps->version = (uint8_t) el[0];
						break;
					case STATE:
						wps->state = (uint8_t) el[0];
						break;
					case LOCKED:
						wps->locked = (uint8_t) el[0];
						break;
					case MANUFACTURER:
						ptr = wps->manufacturer;
						break;
					case MODEL_NAME:
						ptr = wps->model_name;
						break;
					case MODEL_NUMBER:
						ptr = wps->model_number;
						break;
					case DEVICE_NAME:
						ptr = wps->device_name;
						break;
					case SSID:
						ptr = wps->ssid;
						break;
					case UUID:
						src = hex2str(el, el_len);
						ptr = wps->uuid;
						break;
					case SERIAL:
						src = hex2str(el, el_len);
						ptr = wps->serial;
						break;
					case SELECTED_REGISTRAR:
						src = hex2str(el, el_len);
						ptr = wps->selected_registrar;
						break;
					case RESPONSE_TYPE:
						src = hex2str(el, el_len);
						ptr = wps->response_type;
						break;
					case PRIMARY_DEVICE_TYPE:
						src = hex2str(el, el_len);
						ptr = wps->primary_device_type;
						break;
					case CONFIG_METHODS:
						src = hex2str(el, el_len);
						ptr = wps->config_methods;
						break;
					case RF_BANDS:
						src = hex2str(el, el_len);
						ptr = wps->rf_bands;
						break;
					case OS_VERSION:
						src = hex2str(el, el_len);
						ptr = wps->os_version;
						break;
					default:
						src = NULL;
						ptr = NULL; 
				}

				if(!ptr)
				{
					continue;
				}

				memset(ptr, 0, LIBWPS_MAX_STR_LEN);
				if(!src)
				{
					memcpy(ptr, el, el_len);
				}
				else
				{
					strncpy(ptr, src, LIBWPS_MAX_STR_LEN);
					free(src);
				}

				src = NULL;
				ptr = NULL;
				free(el);
			}
		}

		ret_val = 1;
		free(wps_ie_data);
	} 

	return ret_val;
}

/* Extracts and returns the WPS IE from a beacon/probe response packet */
unsigned char *get_wps_data(const u_char *data, size_t len, size_t *tag_len)
{
	unsigned char *tag_data = NULL;
	struct tagged_parameter *tag_params = NULL;
	size_t tag_data_len = 0;
	int i = 0;

	for(i=0; i<len; i++)
	{
		/* Look for the WPS IE tag number */
		if(data[i] != WPS_TAG_NUMBER)
		{
			continue;
		}

		/* Double check remaining size in data buffer */
		if((len - i) > (WPS_VENDOR_ID_SIZE + VENDOR_ID_OFFSET))
		{
			/* 
			 * The WPS IE tag number is 'Vendor Specific' and has various other uses. 
			 * Verify that this is actually a WPS IE.
			 */
			if(memcmp((data+i+VENDOR_ID_OFFSET), WPS_VENDOR_ID, WPS_VENDOR_ID_SIZE) == 0)
			{
				tag_params = (struct tagged_parameter *) (data+i);
				tag_data_len = tag_params->len - WPS_VENDOR_ID_SIZE;

				tag_data = malloc(tag_data_len);
				if(tag_data)
				{
					memset(tag_data, 0, tag_data_len);
					memcpy(tag_data, (data+i+VENDOR_ID_OFFSET+WPS_VENDOR_ID_SIZE), tag_data_len);
					*tag_len = tag_data_len;
				}
	
				break;
			}
		}
	}

        return tag_data;
}

/* Gets the data for a given IE inside a tagged parameter list */
unsigned char *get_wps_data_element(const u_char *data, size_t len, uint16_t type, size_t *el_len)
{
        unsigned char *el_data = NULL;
        int offset = 0, tag_size = 0, el_data_size = 0;
        struct data_element *el = NULL;

        tag_size = sizeof(struct data_element);

        while((offset + tag_size) < len)
        {
                el = (struct data_element *) (data + offset);

                /* Check for the tag number and a sane tag length value */
                if((ntohs(el->type) == type) &&
                   (ntohs(el->len) <= (len - offset - tag_size))
                )
                {
			el_data_size = ntohs(el->len);
                        el_data = malloc(el_data_size);
                        if(el_data)
                        {
                                memset(el_data, 0, el_data_size);
                                memcpy(el_data, (data + offset + tag_size), el_data_size);
                                *el_len = el_data_size;
                        }
                        break;
		}

                offset += (tag_size + ntohs(el->len));
        }

        return el_data;
}

/* Make best guess to determine if a radio tap header is present */
int libwps_has_rt_header(const u_char *packet, size_t len)
{
        struct radio_tap_header *rt_header = 0;
        int yn = 1;

        rt_header = (struct radio_tap_header *) packet;

        if((rt_header->revision != RADIO_TAP_VERSION) ||
           ((int) rt_header->len <= 0) ||
           (rt_header->len >= len)
          )
        {
                yn = 0;
        }

        return yn;
}

/* 
 * Returns a pointer to the radio tap header. If there is no radio tap header,
 * it returns a pointer to a dummy radio tap header.
 */
const u_char *libwps_radio_header(const u_char *packet, size_t len)
{
        if(libwps_has_rt_header(packet, len))
        {
                return packet;
        }
        else
        {
                return (u_char *) FAKE_RADIO_TAP_HEADER;
        }

}

/* Convert raw data to a hex string */
char *hex2str(unsigned char *hex, int len)
{
	static const char atab[] = "0123456789abcdef";
	char *str = malloc((len*2)+1), *out=str;
	if(!str) return 0;
	for(;len;hex++,len--) {
		*(out++) = atab[*hex >> 4];
		*(out++) = atab[*hex & 0xf];
	}
	*out = 0;
	return str;
}
