Merge branch 'hmoore-r7-feature/add-udp-probe-specifiers' of github.com:zmap/zmap into stable

src/probe_modules/module_udp.c

@@ -25,12 +25,111 @@
 
 #include "probe_modules.h"
 #include "packet.h"
+#include "logger.h"
+
+#define MAX_UDP_PAYLOAD_LEN 1472
+
+char *udp_send_msg = NULL; // Must be null-terminated
+int udp_send_msg_len = 0;
+
+const char *udp_send_msg_default = "GET / HTTP/1.1\r\n\r\n"; 
 
-const char *udp_send_msg = "GET / HTTP/1.1\r\n\r\n"; // Must be null-terminated
 static int num_ports = 1;
 
 probe_module_t module_udp;
 
+
+int udp_global_initialize(struct state_conf * zconf) {
+	char *args, *c;
+	int i;
+	unsigned int n;
+
+	FILE *inp;
+
+	udp_send_msg = strdup(udp_send_msg_default);
+	udp_send_msg_len = strlen(udp_send_msg);
+
+	if (! (zconf->probe_args && strlen(zconf->probe_args) > 0)) 
+		return(0);
+	
+	args = strdup(zconf->probe_args);
+	if (! args) exit(1);
+
+	c = strchr(args, ':');
+	if (! c) {
+		free(args);
+		free(udp_send_msg);
+		log_fatal("udp", "unknown UDP probe specification (expected file:/path or text:STRING or hex:01020304)");
+		exit(1);
+	}
+
+	*c++ = 0;
+
+	if (strcmp(args, "text") == 0) {
+		udp_send_msg = strdup(c);
+		udp_send_msg_len = strlen(udp_send_msg);
+
+	} else if (strcmp(args, "file") == 0) {
+		inp = fopen(c, "rb");
+		if (!inp) {
+			free(args);
+			free(udp_send_msg);
+			log_fatal("udp", "could not open UDP data file '%s'\n", c);
+			exit(1);
+		}
+		udp_send_msg = malloc(MAX_UDP_PAYLOAD_LEN);
+		if (! udp_send_msg) {
+			free(args);
+			free(udp_send_msg);
+			log_fatal("udp", "failed to malloc payload buffer");
+			exit(1);
+		}
+		udp_send_msg_len = fread(udp_send_msg, 1, MAX_UDP_PAYLOAD_LEN, inp);
+		fclose(inp);
+
+	} else if (strcmp(args, "hex") == 0) {
+		udp_send_msg_len = strlen(c) / 2;
+		udp_send_msg = malloc(udp_send_msg_len);
+		if (! udp_send_msg) {
+			free(args);
+			free(udp_send_msg);
+			log_fatal("udp", "failed to malloc payload buffer");
+			exit(1);
+		}
+
+		for (i=0; i < udp_send_msg_len; i++) {
+			if (sscanf(c + (i*2), "%2x", &n) != 1) {
+				free(args);
+				free(udp_send_msg);
+				log_fatal("udp", "non-hex character: '%c'", c[i*2]);
+				exit(1);
+			}
+			udp_send_msg[i] = (n & 0xff);
+		}
+	} else {
+		log_fatal("udp", "unknown UDP probe specification (expected file:/path, text:STRING, or hex:01020304)");
+		free(args);
+		exit(1);
+	}
+
+	if (udp_send_msg_len > MAX_UDP_PAYLOAD_LEN) {
+		fprintf(stderr, "warning: reducing UDP payload to %d bytes (from %d) to fit on the wire\n", 
+				MAX_UDP_PAYLOAD_LEN, udp_send_msg_len);
+		udp_send_msg_len = MAX_UDP_PAYLOAD_LEN;
+	}
+	free(args);
+	return(0);
+}
+
+int udp_global_cleanup(__attribute__((unused)) struct state_conf *zconf,
+					   __attribute__((unused)) struct state_send *zsend,
+					   __attribute__((unused)) struct state_recv *zrecv) {
+	if (udp_send_msg) free(udp_send_msg);
+	udp_send_msg = NULL;
+	return(0);
+}
+
+
 int udp_init_perthread(void* buf, macaddr_t *src,
 		macaddr_t *gw, __attribute__((unused)) port_h_t dst_port)
 {
@@ -38,20 +137,20 @@ int udp_init_perthread(void* buf, macaddr_t *src,
 	struct ethhdr *eth_header = (struct ethhdr *)buf;
 	make_eth_header(eth_header, src, gw);
 	struct iphdr *ip_header = (struct iphdr*)(&eth_header[1]);
-	uint16_t len = htons(sizeof(struct iphdr) + sizeof(struct udphdr) + strlen(udp_send_msg));
+	uint16_t len = htons(sizeof(struct iphdr) + sizeof(struct udphdr) + udp_send_msg_len);
 	make_ip_header(ip_header, IPPROTO_UDP, len);
 
 	struct udphdr *udp_header = (struct udphdr*)(&ip_header[1]);
-	len = sizeof(struct udphdr) + strlen(udp_send_msg);
+	len = sizeof(struct udphdr) + udp_send_msg_len;
 	make_udp_header(udp_header, zconf.target_port, len);
 
 	char* payload = (char*)(&udp_header[1]);
 
 	module_udp.packet_length = sizeof(struct ethhdr) + sizeof(struct iphdr) 
-				+ sizeof(struct udphdr) + strlen(udp_send_msg);
+				+ sizeof(struct udphdr) + udp_send_msg_len;
 	assert(module_udp.packet_length <= MAX_PACKET_SIZE);
 
-	strcpy(payload, udp_send_msg);
+	memcpy(payload, udp_send_msg, udp_send_msg_len);
 
 	num_ports = zconf.source_port_last - zconf.source_port_first + 1;
 
@@ -212,16 +311,17 @@ static response_type_t responses[] = {
 
 probe_module_t module_udp = {
 	.name = "udp",
-	.packet_length = 96,
+	.packet_length = 1,
 	.pcap_filter = "udp || icmp",
-	.pcap_snaplen = 96,
+	.pcap_snaplen = 1500,
 	.port_args = 1,
 	.thread_initialize = &udp_init_perthread,
+	.global_initialize = &udp_global_initialize,
 	.make_packet = &udp_make_packet,
 	.print_packet = &udp_print_packet,
 	.validate_packet = &udp_validate_packet,
 	.classify_packet = &udp_classify_packet,
-	.close = NULL,
+	.close = &udp_global_cleanup,
 	.responses = responses
 };