/* This file is Copyright (c) 2002 Gianni Tedesco.
 * This code is released under the terms of the GNU GPL version 2
 *
 * This is a free implementation of the WISP protocol as used by
 * websense for interoperating with proxies. Currently there are
 * no free implementations that I am aware of (or even non-free
 * implementations for squid).
 *
 * Basically you send the server a URL and it says yes or no.
 */

#include <stdlib.h>
#include <stdio.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <string.h>
#include <netdb.h>
#include <unistd.h>

/* u?? fields are unknown to us right now */

#define WISP_MAGIC 0x00810000

struct wisp_req {
	u_int32_t u1; /* 0x00300001 */
	u_int32_t magic; 
	u_int32_t serial; /* unique each time */
	u_int32_t u2; /* 0x00010000 */
	u_int32_t saddr; /* IP address of client doing request */
	u_int32_t daddr; /* IP address of server */
	u_int32_t len; /* Length of following URL in bytes */
	/* URL */
};

struct wisp_resp {
	u_int32_t u1; 
	u_int32_t magic; /* 0x00810000 */
	u_int32_t serial; /* echoed from request */
	u_int32_t u2;
	u_int32_t len; /* Length of following HTTP resoponse in bytes */
	/* HTTP 1.0 reply header */
};

static char buf[8192];

/* Initialse a websense query with a URL */
int wisp_query(char *buf, int buflen, char *url)
{
	struct wisp_req *req=(struct wisp_req *)buf;
	int url_len;

	if ( !buf || ! url ) return -1;

	/* Calculate length */
	url_len=strlen(url);
	if ( buflen < sizeof(*req)+url_len ) return -1;
	
	/* Fill in the packet fields */
	req->u1=htonl(0x00300001);
	req->magic=htonl(WISP_MAGIC);
	req->serial=htonl(0xdeadbeef);
	req->u2=htonl(0x00010000);
	req->len=htonl(url_len);

	/* Copy in the header directly followed by the URL */
	memcpy(buf, (char *)req, sizeof(*req));
	memcpy(buf+sizeof(*req), url, url_len);

	return sizeof(*req) + url_len;
}

/* Set the addresses for a websense query */
void wisp_addr(char *buf, u_int32_t s, u_int32_t d)
{
	struct wisp_req *req=(struct wisp_req *)buf;

	if ( !req ) return;

	req->saddr=s;
	req->daddr=d;
}

/* Produce addresses suitable for wisp_addr from a string */
u_int32_t wisp_resolve(char *addr)
{
	struct hostent *h;
	struct sockaddr_in sin;

	if ( !(h=gethostbyname(addr)) ) {
		perror("gethostbyname()");
		exit(1);
	}

	if ( h->h_addrtype != AF_INET ) {
		printf("Must use IPv4 addresses (%s)\n", addr);
		exit(1);
	}

	memcpy(&sin.sin_addr, h->h_addr, sizeof(sin.sin_addr));
	return sin.sin_addr.s_addr;
}

/* 4 args: wsvr src dst url */
int main(int argc, char **argv)
{
	int fd;
	int rlen;
	int req_len;
	struct wisp_resp *reply;
	struct hostent *h;
	struct sockaddr_in sin;

	u_int32_t src_addr, dst_addr;

	if ( argc < 5 ) {
		printf("%s: websense_ip srcip dstip url\n", argv[0]);
		exit(1);
	}

	/* We tell webschlong the request came from this one */
	src_addr=wisp_resolve(argv[2]);
	dst_addr=wisp_resolve(argv[3]);

	/* We connect to this one */
	if ( !(h=gethostbyname(argv[1])) ) {
		perror("gethostbyname()");
		exit(1);
	}
	memcpy(&sin.sin_addr, h->h_addr, sizeof(sin.sin_addr));

	/* Create a socket */
	if ( (fd=socket(AF_INET, SOCK_STREAM, 0))<0 ) {
		perror("socket()");
		exit(1);
	}

	/* Connect to websense */
	sin.sin_family=AF_INET;
	sin.sin_port=htons(15868);
	if ( connect(fd, (struct sockaddr *)&sin, sizeof(sin))<0 ) {
		perror("connect()");
		exit(1);
	}

	/* Send the query */
	req_len=wisp_query(buf, sizeof(buf), argv[4]);
	wisp_addr(buf, src_addr, dst_addr);
	send(fd, buf, req_len, 0);

	/* Wait for a response */
	if ( (rlen=recv(fd, buf, sizeof(buf), 0))<0 ) {
		printf("read()");
		exit(1);
	}

	/* Check we got at least the minimal structure back */
	reply=(struct wisp_resp *)buf;
	if ( rlen<12 ) {
		printf("server returned short packet (%i bytes)\n", rlen);
		exit(1);
	}

	/* Print out unknown flags */
	printf("unknown1: 0x%.8x:\n", ntohl(reply->u1));
	if ( rlen >= 16 ) 
		printf("unknown2: 0x%.8x:\n", ntohl(reply->u2));
	printf("\n");

	if ( rlen >= 20 ) {
		int len=ntohl(reply->len);
		if ( len > rlen-sizeof(*reply) ) len=rlen-sizeof(reply);
		write(1, buf+16, len);
		write(1, "\n", 1);
	}else{
		printf("URL is OK, I guess...\n");
	}

	close(fd);
	exit(0);
}