#include "check_upnp.h"

using namespace std;

char *servicename = (char*)"UPnP";

void printVersion()
{
	cout << "check_upnp v" << VERSION << endl << endl;
}

void printHelp(bool longVersion)
{
	if(longVersion)
	{
		printVersion();
		cout << "Check UPnP server" << endl << endl;
		printHelp(false);
		cout << "Options:" << endl;
		cout << " -h" << endl;
		cout << "    Print detailed help screen" << endl;
		cout << " -V" << endl;
		cout << "    Print version information" << endl;
		cout << " -H HOSTADDRESS" << endl;
		cout << "    Host where the UPnP server is running" << endl;
		return;
	}
	cout << "Usage: " << endl << "check_upnp [-hV] -H HOSTADDRESS" << endl << endl;
}

int check_upnp(char *hostname, string *serverinfo)
{
    char buffer[MAX_UDP+1];
    struct sockaddr_in si;
    struct hostent *host = NULL;
    int s;
    int timeout = 10;
    struct hostent *he;
    struct in_addr **addr_list;

    he = gethostbyname(hostname);
    if(he == NULL)
    {
        *serverinfo = "Can't resolve " + string(hostname);
        return 3;
    }
    addr_list = (struct in_addr **) he->h_addr_list;
    char targetIP[16];
    strcpy(targetIP,inet_ntoa(*addr_list[0]));

    setupSocket(1900,(char*)"239.255.255.250",host,timeout,&si,&s);
    memset(buffer,0x00,MAX_UDP+1);
	char *discover = (char*) "M-SEARCH * HTTP/1.1\r\nHOST: 239.255.255.250:1900\r\nMAN: \"ssdp:discover\"\r\nMX: 5\r\nST: urn:schemas-upnp-org:device:MediaServer:1\r\n\r\n";
	sendMsg(s,discover,strlen(discover),&si);
    int recvstatus;
    while(true)
    {
        recvstatus = recvMsg(s, buffer, MAX_UDP, &si);
        if(recvstatus == -1)
        {
            *serverinfo = "No response";
            return 2;
        }

        char *sourceIP = inet_ntoa(si.sin_addr);
        if(!strcmp(sourceIP,targetIP))
        {
            stringstream ss(buffer);
            string line;
            while(getline(ss,line,'\n'))
            {
                if(line.length() > strlen("SERVER:"))
                {
                    if (line.substr(0,strlen("SERVER:")) == "SERVER:")
                    {
                        *serverinfo = line.substr(0,line.size() - 1);
                    }
                }
            }
            return 0;
        }

    }
}

int main(int argc, char **argv)
{
	struct itimerval timer;
	timer.it_value.tv_sec = 10;
	timer.it_value.tv_usec = 0;
	timer.it_interval.tv_sec = 0;
	timer.it_interval.tv_usec = 0;
	setitimer (ITIMER_VIRTUAL, &timer, 0);

	struct sigaction sa;
	memset (&sa, 0, sizeof (sa));
	sa.sa_handler = &timer_handler;
	sigaction (SIGVTALRM, &sa, 0);

	uint16_t port = 27015;
	char *hostname = NULL;
	int c;

	while ((c = getopt (argc, argv, "H:p:Vh")) != -1)
	{
		switch(c)
		{
			case 'H':
        			hostname = optarg;
        			break;
			case 'p':
       		 		port = (uint16_t) str2int(optarg);
       		 		break;
			case 'V':
				printVersion();
				return 0;
			case 'h':
				printHelp(true);
				return 0;
			case '?':
				printHelp(false);
				return 3;
		}
	}

	if(hostname == NULL)
	{
		cout << "No HOSTADDRESS specified. Exiting." << endl;
		return 3;
	}

	string serverinfo = "";
	int returnCode = check_upnp(hostname,&serverinfo);

	cout << servicename;
	switch(returnCode)
	{
		case 0:
			cout << " OK - " << serverinfo << endl ;
			break;

		case 2:
			cout << " CRITICAL - No response" << endl;
			break;

        case 3:
            cout << " UNKNOWN - " << serverinfo << endl;
	}

	return returnCode;
}