1/*
2 * Original implementation on libmnl:
3 * (C) 2011 by Pablo Neira Ayuso <pablo@netfilter.org>
4 * (C) 2011 by Intra2net AG <http://www.intra2net.com>
5 *
6 * Port to libnl:
7 * (C) 2013 by Mathieu J. Poirier <mathieu.poirier@linaro.org>
8 *
9 * This program is free software; you can redistribute it and/or modify
10 * it under the terms of the GNU General Public License as published
11 * by the Free Software Foundation; either version 2 of the License, or
12 * (at your option) any later version.
13 */
14
15#include <stdio.h>
16#include <stdlib.h>
17#include <string.h>
18#include <sys/types.h>
19#include <sys/socket.h>
20#include <dirent.h>
21#include <sys/stat.h>
22#include <fcntl.h>
23#include <unistd.h>
24#include <time.h>
25#include <errno.h>
26
27#include <netlink-local.h>
28#include <linux/netlink.h>
29#include <linux/netfilter/nfnetlink.h>
30#include <linux/netfilter/nfnetlink_acct.h>
31#include <netlink/netfilter/nfnl.h>
32#include <netlink/netlink.h>
33#include <netlink/socket.h>
34#include <netlink/msg.h>
35
36#define VERSION "1.0.1"
37
38enum {
39	NFACCT_CMD_NONE = 0,
40	NFACCT_CMD_LIST,
41	NFACCT_CMD_ADD,
42	NFACCT_CMD_DELETE,
43	NFACCT_CMD_GET,
44	NFACCT_CMD_FLUSH,
45	NFACCT_CMD_VERSION,
46	NFACCT_CMD_HELP,
47	NFACCT_CMD_RESTORE,
48};
49
50static int nfacct_cmd_list(int argc, char *argv[]);
51static int nfacct_cmd_add(int argc, char *argv[]);
52static int nfacct_cmd_delete(int argc, char *argv[]);
53static int nfacct_cmd_get(int argc, char *argv[]);
54static int nfacct_cmd_flush(int argc, char *argv[]);
55static int nfacct_cmd_version(int argc, char *argv[]);
56static int nfacct_cmd_help(int argc, char *argv[]);
57static int nfacct_cmd_restore(int argc, char *argv[]);
58
59#ifndef HAVE_LIBNL20
60#define nl_sock nl_handle
61#define nl_socket_alloc nl_handle_alloc
62#define nl_socket_free nl_handle_destroy
63#endif
64
65static void usage(char *argv[])
66{
67	fprintf(stderr, "Usage: %s command [parameters]...\n", argv[0]);
68}
69
70static void nfacct_perror(const char *msg)
71{
72	if (errno == 0) {
73		fprintf(stderr, "nfacct v%s: %s\n", VERSION, msg);
74	} else {
75		fprintf(stderr, "nfacct v%s: %s: %s\n",
76			VERSION, msg, strerror(errno));
77	}
78}
79
80int main(int argc, char *argv[])
81{
82	int cmd = NFACCT_CMD_NONE, ret = 0;
83
84	if (argc < 2) {
85		usage(argv);
86		exit(EXIT_FAILURE);
87	}
88
89	if (strncmp(argv[1], "list", strlen(argv[1])) == 0)
90		cmd = NFACCT_CMD_LIST;
91	else if (strncmp(argv[1], "add", strlen(argv[1])) == 0)
92		cmd = NFACCT_CMD_ADD;
93	else if (strncmp(argv[1], "delete", strlen(argv[1])) == 0)
94		cmd = NFACCT_CMD_DELETE;
95	else if (strncmp(argv[1], "get", strlen(argv[1])) == 0)
96		cmd = NFACCT_CMD_GET;
97	else if (strncmp(argv[1], "flush", strlen(argv[1])) == 0)
98		cmd = NFACCT_CMD_FLUSH;
99	else if (strncmp(argv[1], "version", strlen(argv[1])) == 0)
100		cmd = NFACCT_CMD_VERSION;
101	else if (strncmp(argv[1], "help", strlen(argv[1])) == 0)
102		cmd = NFACCT_CMD_HELP;
103	else if (strncmp(argv[1], "restore", strlen(argv[1])) == 0)
104		cmd = NFACCT_CMD_RESTORE;
105	else {
106		fprintf(stderr, "nfacct v%s: Unknown command: %s\n",
107			VERSION, argv[1]);
108		usage(argv);
109		exit(EXIT_FAILURE);
110	}
111
112	switch(cmd) {
113	case NFACCT_CMD_LIST:
114		ret = nfacct_cmd_list(argc, argv);
115		break;
116	case NFACCT_CMD_ADD:
117		ret = nfacct_cmd_add(argc, argv);
118		break;
119	case NFACCT_CMD_DELETE:
120		ret = nfacct_cmd_delete(argc, argv);
121		break;
122	case NFACCT_CMD_GET:
123		ret = nfacct_cmd_get(argc, argv);
124		break;
125	case NFACCT_CMD_FLUSH:
126		ret = nfacct_cmd_flush(argc, argv);
127		break;
128	case NFACCT_CMD_VERSION:
129		ret = nfacct_cmd_version(argc, argv);
130		break;
131	case NFACCT_CMD_HELP:
132		ret = nfacct_cmd_help(argc, argv);
133		break;
134	case NFACCT_CMD_RESTORE:
135		ret = nfacct_cmd_restore(argc, argv);
136		break;
137	}
138	return ret < 0 ? EXIT_FAILURE : EXIT_SUCCESS;
139}
140
141
142static int message_received(struct nl_msg *msg, void *arg)
143{
144	struct nlmsghdr *hdr = msg->nm_nlh;
145
146	if (hdr->nlmsg_type == NLMSG_ERROR) {
147		struct nlmsgerr *err = nlmsg_data(hdr);
148
149		if (err->error == 0)
150			return NL_STOP;
151	}
152
153	return NL_OK;
154}
155
156static int valid_input(struct nl_msg *msg, void *arg)
157{
158	struct nlattr *nla = nlmsg_attrdata(nlmsg_hdr(msg),
159					 sizeof(struct nfgenmsg));
160	struct nlattr *tb[NFACCT_NAME_MAX+1] = {};
161	char buf[4096];
162	int ret;
163
164	ret = nlmsg_parse(nlmsg_hdr(msg),
165			 sizeof(struct nfgenmsg), tb, NFACCT_MAX, NULL);
166
167	if (ret < 0) {
168		nfacct_perror("Can't parse message\n");
169		return ret;
170	}
171
172	ret = snprintf(buf, sizeof(buf),
173		"{ pkts = %.20llu, bytes = %.20llu } = %s;",
174		(unsigned long long)be64toh(nla_get_u64(tb[NFACCT_PKTS])),
175		(unsigned long long)be64toh(nla_get_u64(tb[NFACCT_BYTES])),
176		nla_get_string(tb[NFACCT_NAME]));
177
178	printf("%s\n", buf);
179
180	return 0;
181}
182
183static int nfacct_cmd_list(int argc, char *argv[])
184{
185	struct nl_msg *msg;
186	struct nl_sock *handle;
187	int zeroctr = 0;
188	int ret, i;
189
190	for (i=2; i<argc; i++) {
191		if (strncmp(argv[i], "reset", strlen(argv[i])) == 0) {
192			zeroctr = 1;
193		} else if (strncmp(argv[i], "xml", strlen(argv[i])) == 0) {
194			nfacct_perror("xml feature not implemented");
195			return -1;
196		} else {
197			nfacct_perror("unknown argument");
198			return -1;
199		}
200	}
201
202	msg = nlmsg_alloc();
203	if (!msg)
204		return -1;
205
206	ret = nfnlmsg_put(msg,
207			NL_AUTO_PID,
208			NL_AUTO_SEQ,
209			NFNL_SUBSYS_ACCT,
210			zeroctr ?
211			NFNL_MSG_ACCT_GET_CTRZERO : NFNL_MSG_ACCT_GET,
212			NLM_F_DUMP | NLM_F_REQUEST,
213			AF_UNSPEC,
214			0);
215
216	if (ret) {
217		NL_DBG(2, "Can't append payload to message: %s line: %d\n",
218							__FUNCTION__, __LINE__);
219		goto fail;
220	}
221
222	handle = nl_socket_alloc();
223	if ((ret = nfnl_connect(handle))) {
224		NL_DBG(2, "Can't connect handle: %s line: %d\n",
225							__FUNCTION__, __LINE__);
226		goto fail;
227	}
228
229	if ((ret = nl_send_auto_complete(handle, msg)) < 0) {
230		NL_DBG(2, "Can't send msg: %s line: %d\n",
231							__FUNCTION__, __LINE__);
232		goto fail_send;
233        }
234
235	nl_socket_modify_cb(handle, NL_CB_VALID, NL_CB_CUSTOM, valid_input, NULL);
236	ret = nl_recvmsgs_default(handle);
237	if (ret < 0) {
238		NL_DBG(2, "Can't receice msg: %s line: %d\n",
239							__FUNCTION__, __LINE__);
240	}
241
242fail_send:
243	nl_close(handle);
244	nl_socket_free(handle);
245fail:
246	nlmsg_free(msg);
247	return ret;
248}
249
250static int _nfacct_cmd_add(char *name, int pkts, int bytes)
251{
252	struct nl_msg *msg;
253	struct nl_sock *handle;
254	char nfname[NFACCT_NAME_MAX];
255	int ret;
256
257	strncpy(nfname, name, NFACCT_NAME_MAX);
258	nfname[NFACCT_NAME_MAX-1] = '\0';
259
260	msg = nlmsg_alloc();
261	if (!msg)
262		return -1;
263
264	ret = nfnlmsg_put(msg,
265			NL_AUTO_PID,
266			NL_AUTO_SEQ,
267			NFNL_SUBSYS_ACCT,
268			NFNL_MSG_ACCT_NEW,
269			NLM_F_CREATE | NLM_F_ACK | NLM_F_REQUEST,
270			AF_UNSPEC,
271			0);
272
273	if (ret) {
274		NL_DBG(2, "Can't append payload to message: %s line: %d\n",
275							__FUNCTION__, __LINE__);
276		goto fail;
277	}
278
279	nla_put_string(msg, NFACCT_NAME, nfname);
280	nla_put_u64(msg, NFACCT_PKTS, htobe64(pkts));
281	nla_put_u64(msg, NFACCT_BYTES, htobe64(bytes));
282
283	handle = nl_socket_alloc();
284	if ((ret = nfnl_connect(handle))) {
285		NL_DBG(2, "Can't connect handle: %s line: %d\n",
286							__FUNCTION__, __LINE__);
287		goto fail;
288	}
289
290	if ((ret = nl_send_auto_complete(handle, msg)) < 0) {
291		NL_DBG(2, "Can't send msg: %s line: %d\n",
292							__FUNCTION__, __LINE__);
293		goto fail_send;
294        }
295
296	ret = nl_recvmsgs_default(handle);
297	if (ret < 0) {
298		NL_DBG(2, "Can't receice msg: %s line: %d\n",
299							__FUNCTION__, __LINE__);
300	}
301
302fail_send:
303	nl_close(handle);
304	nl_socket_free(handle);
305fail:
306	nlmsg_free(msg);
307	return ret;
308}
309
310
311
312static int nfacct_cmd_add(int argc, char *argv[])
313{
314	if (argc < 3) {
315		nfacct_perror("missing object name");
316		return -1;
317	} else if (argc > 3) {
318		nfacct_perror("too many arguments");
319		return -1;
320	}
321
322	return _nfacct_cmd_add(argv[2], 0, 0);
323}
324
325static int nfacct_cmd_delete(int argc, char *argv[])
326{
327	struct nl_msg *msg;
328	struct nl_sock *handle;
329	char nfname[NFACCT_NAME_MAX];
330	int ret;
331
332	if (argc < 3) {
333		nfacct_perror("missing object name");
334		return -1;
335	} else if (argc > 3) {
336		nfacct_perror("too many arguments");
337		return -1;
338	}
339
340	strncpy(nfname, argv[2], NFACCT_NAME_MAX);
341	nfname[NFACCT_NAME_MAX-1] = '\0';
342
343	msg = nlmsg_alloc();
344	if (!msg)
345		return -1;
346
347	ret = nfnlmsg_put(msg,
348			NL_AUTO_PID,
349			NL_AUTO_SEQ,
350			NFNL_SUBSYS_ACCT,
351			NFNL_MSG_ACCT_DEL,
352			NLM_F_ACK | NLM_F_REQUEST,
353			AF_UNSPEC,
354			0);
355
356	if (ret) {
357		NL_DBG(2, "Can't append payload to message: %s line: %d\n",
358							__FUNCTION__, __LINE__);
359		goto fail;
360	}
361
362	nla_put_string(msg, NFACCT_NAME, nfname);
363
364	handle = nl_socket_alloc();
365	if ((ret = nfnl_connect(handle))) {
366		NL_DBG(2, "Can't connect handle: %s line: %d\n",
367							__FUNCTION__, __LINE__);
368		goto fail;
369	}
370
371	if ((ret = nl_send_auto_complete(handle, msg)) < 0) {
372		NL_DBG(2, "Can't send msg: %s line: %d\n",
373							__FUNCTION__, __LINE__);
374		goto fail_send;
375        }
376
377	ret = nl_recvmsgs_default(handle);
378	if (ret < 0) {
379		NL_DBG(2, "Can't receice msg: %s line: %d\n",
380							__FUNCTION__, __LINE__);
381	}
382
383fail_send:
384	nl_close(handle);
385	nl_socket_free(handle);
386fail:
387	nlmsg_free(msg);
388	return ret;
389	return 0;
390}
391
392
393static int nfacct_cmd_get(int argc, char *argv[])
394{
395	struct nl_msg *msg;
396	struct nl_sock *handle;
397	struct nl_cb *cb;
398	char nfname[NFACCT_NAME_MAX];
399	int zeroctr = 0;
400	int ret, i;
401
402	if (argc < 3) {
403		nfacct_perror("missing object name");
404		 return -1;
405	}
406
407	for (i=3; i<argc; i++) {
408		if (strncmp(argv[i], "reset", strlen(argv[i])) == 0) {
409			zeroctr = 1;
410		} else if (strncmp(argv[i], "xml", strlen(argv[i])) == 0) {
411			nfacct_perror("xml feature not implemented");
412			return -1;
413		} else {
414			nfacct_perror("unknown argument");
415			return -1;
416		}
417	}
418
419	strncpy(nfname, argv[2], NFACCT_NAME_MAX);
420	nfname[NFACCT_NAME_MAX-1] = '\0';
421
422	msg = nlmsg_alloc();
423	if (!msg)
424		return -1;
425
426	ret = nfnlmsg_put(msg,
427			NL_AUTO_PID,
428			NL_AUTO_SEQ,
429			NFNL_SUBSYS_ACCT,
430			zeroctr ?
431			NFNL_MSG_ACCT_GET_CTRZERO : NFNL_MSG_ACCT_GET,
432			NLM_F_ACK | NLM_F_REQUEST,
433			AF_UNSPEC,
434			0);
435
436	if (ret) {
437		NL_DBG(2, "Can't append payload to message: %s line: %d\n",
438							__FUNCTION__, __LINE__);
439		goto fail;
440	}
441
442	nla_put_string(msg, NFACCT_NAME, nfname);
443
444	handle = nl_socket_alloc();
445
446	if (handle) {
447		cb = nl_cb_alloc(NL_CB_DEFAULT);
448		if (!cb)
449			goto fail;
450
451		if (nl_cb_set(cb, NL_CB_MSG_IN,
452				 NL_CB_CUSTOM,
453				 message_received, NULL) < 0)
454			goto fail;
455
456		nl_socket_set_cb(handle,cb);
457	} else {
458		goto fail;
459	}
460
461	if ((ret = nfnl_connect(handle))) {
462		NL_DBG(2, "Can't connect handle: %s line: %d\n",
463							__FUNCTION__, __LINE__);
464		goto fail;
465	}
466
467	if ((ret = nl_send_auto_complete(handle, msg)) < 0) {
468		NL_DBG(2, "Can't send msg: %s line: %d\n",
469							__FUNCTION__, __LINE__);
470		goto fail_send;
471        }
472
473	nl_socket_modify_cb(handle, NL_CB_VALID, NL_CB_CUSTOM, valid_input, NULL);
474	ret = nl_recvmsgs_default(handle);
475	if (ret < 0) {
476		NL_DBG(2, "Can't receice msg: %s line: %d\n",
477							__FUNCTION__, __LINE__);
478	}
479
480fail_send:
481	nl_close(handle);
482	nl_socket_free(handle);
483fail:
484	nlmsg_free(msg);
485	return ret;
486}
487
488static int nfacct_cmd_flush(int argc, char *argv[])
489{
490	struct nl_msg *msg;
491	struct nl_sock *handle;
492	int ret;
493
494	if (argc > 2) {
495		nfacct_perror("too many arguments");
496		return -1;
497	}
498
499	msg = nlmsg_alloc();
500	if (!msg)
501		return -1;
502
503	ret = nfnlmsg_put(msg,
504			NL_AUTO_PID,
505			NL_AUTO_SEQ,
506			NFNL_SUBSYS_ACCT,
507			NFNL_MSG_ACCT_DEL,
508			NLM_F_ACK | NLM_F_REQUEST,
509			AF_UNSPEC,
510			0);
511
512	if (ret) {
513		NL_DBG(2, "Can't append payload to message: %s line: %d\n",
514							__FUNCTION__, __LINE__);
515		goto fail;
516	}
517
518	handle = nl_socket_alloc();
519	if ((ret = nfnl_connect(handle))) {
520		NL_DBG(2, "Can't connect handle: %s line: %d\n",
521							__FUNCTION__, __LINE__);
522		goto fail;
523	}
524
525	if ((ret = nl_send_auto_complete(handle, msg)) < 0) {
526		NL_DBG(2, "Can't send msg: %s line: %d\n",
527							__FUNCTION__, __LINE__);
528		goto fail_send;
529        }
530
531	ret = nl_recvmsgs_default(handle);
532	if (ret < 0) {
533		NL_DBG(2, "Can't receice msg: %s line: %d\n",
534							__FUNCTION__, __LINE__);
535	}
536
537fail_send:
538	nl_close(handle);
539	nl_socket_free(handle);
540fail:
541	nlmsg_free(msg);
542	return ret;
543}
544
545static const char version_msg[] =
546	"nfacct v%s: utility for the Netfilter extended accounting "
547	"infrastructure\n"
548	"Copyright (C) 2011 Pablo Neira Ayuso <pablo@netfilter.org>\n"
549	"Copyright (C) 2011 Intra2net AG <http://www.intra2net.com>\n"
550	"Copyright (C) 2013 Mathieu Poirier <mathieu.poirier@linaro.org>\n"
551	"This program comes with ABSOLUTELY NO WARRANTY.\n"
552	"This is free software, and you are welcome to redistribute it under "
553	"certain \nconditions; see LICENSE file distributed in this package "
554	"for details.\n";
555
556static int nfacct_cmd_version(int argc, char *argv[])
557{
558	printf(version_msg, VERSION);
559	return 0;
560}
561
562static const char help_msg[] =
563	"nfacct v%s: utility for the Netfilter extended accounting "
564	"infrastructure\n"
565	"Usage: %s command [parameters]...\n\n"
566	"Commands:\n"
567	"  list [reset]\t\tList the accounting object table (and reset)\n"
568	"  add object-name\tAdd new accounting object to table\n"
569	"  delete object-name\tDelete existing accounting object\n"
570	"  get object-name\tGet existing accounting object\n"
571	"  flush\t\t\tFlush accounting object table\n"
572	"  restore\t\tRestore accounting object table reading 'list' output from stdin\n"
573	"  version\t\tDisplay version and disclaimer\n"
574	"  help\t\t\tDisplay this help message\n";
575
576static int nfacct_cmd_help(int argc, char *argv[])
577{
578	printf(help_msg, VERSION, argv[0]);
579	return 0;
580}
581
582static int nfacct_cmd_restore(int argc, char *argv[])
583{
584	uint64_t pkts, bytes;
585	char name[512];
586	char buffer[512];
587	int ret;
588	while (fgets(buffer, sizeof(buffer), stdin)) {
589		char *semicolon = strchr(buffer, ';');
590		if (semicolon == NULL) {
591			nfacct_perror("invalid line");
592			return -1;
593		}
594		*semicolon = 0;
595		ret = sscanf(buffer, "{ pkts = %llu, bytes = %llu } = %s",
596		       &pkts, &bytes, name);
597		if (ret != 3) {
598			nfacct_perror("error reading input");
599			return -1;
600		}
601		if ((ret = _nfacct_cmd_add(name, pkts, bytes)) != 0)
602			return ret;
603
604	}
605	return 0;
606}
607