/* ** Copyright (C) 2001-2007 by Carnegie Mellon University. ** ** @OPENSOURCE_HEADER_START@ ** ** Use of the SILK system and related source code is subject to the terms ** of the following licenses: ** ** GNU Public License (GPL) Rights pursuant to Version 2, June 1991 ** Government Purpose License Rights (GPLR) pursuant to DFARS 252.225-7013 ** ** NO WARRANTY ** ** ANY INFORMATION, MATERIALS, SERVICES, INTELLECTUAL PROPERTY OR OTHER ** PROPERTY OR RIGHTS GRANTED OR PROVIDED BY CARNEGIE MELLON UNIVERSITY ** PURSUANT TO THIS LICENSE (HEREINAFTER THE "DELIVERABLES") ARE ON AN ** "AS-IS" BASIS. CARNEGIE MELLON UNIVERSITY MAKES NO WARRANTIES OF ANY ** KIND, EITHER EXPRESS OR IMPLIED AS TO ANY MATTER INCLUDING, BUT NOT ** LIMITED TO, WARRANTY OF FITNESS FOR A PARTICULAR PURPOSE, ** MERCHANTABILITY, INFORMATIONAL CONTENT, NONINFRINGEMENT, OR ERROR-FREE ** OPERATION. CARNEGIE MELLON UNIVERSITY SHALL NOT BE LIABLE FOR INDIRECT, ** SPECIAL OR CONSEQUENTIAL DAMAGES, SUCH AS LOSS OF PROFITS OR INABILITY ** TO USE SAID INTELLECTUAL PROPERTY, UNDER THIS LICENSE, REGARDLESS OF ** WHETHER SUCH PARTY WAS AWARE OF THE POSSIBILITY OF SUCH DAMAGES. ** LICENSEE AGREES THAT IT WILL NOT MAKE ANY WARRANTY ON BEHALF OF ** CARNEGIE MELLON UNIVERSITY, EXPRESS OR IMPLIED, TO ANY PERSON ** CONCERNING THE APPLICATION OF OR THE RESULTS TO BE OBTAINED WITH THE ** DELIVERABLES UNDER THIS LICENSE. ** ** Licensee hereby agrees to defend, indemnify, and hold harmless Carnegie ** Mellon University, its trustees, officers, employees, and agents from ** all claims or demands made against them (and any related losses, ** expenses, or attorney's fees) arising out of, or relating to Licensee's ** and/or its sub licensees' negligent use or willful misuse of or ** negligent conduct or willful misconduct regarding the Software, ** facilities, or other rights or assistance granted by Carnegie Mellon ** University under this License, including, but not limited to, any ** claims of product liability, personal injury, death, damage to ** property, or violation of any laws or regulations. ** ** Carnegie Mellon University Software Engineering Institute authored ** documents are sponsored by the U.S. Department of Defense under ** Contract F19628-00-C-0003. Carnegie Mellon University retains ** copyrights in all material produced under this contract. The U.S. ** Government retains a non-exclusive, royalty-free license to publish or ** reproduce these documents, or allow others to do so, for U.S. ** Government purposes only pursuant to the copyright license under the ** contract clause at 252.227.7013. ** ** @OPENSOURCE_HEADER_END@ */ /* * * rwsettool.c * * FIXME */ #include "silk.h" RCSIDENT("$SiLK: rwsettool.c 7318 2007-05-29 13:57:52Z mthomas $"); #include "iptree.h" #include "utils.h" #include "sksite.h" /* LOCAL DEFINES AND TYPEDEFS */ /* where to write output from --help */ #define USAGE_FH stdout /* LOCAL FUNCTIONS */ static void appUsageLong(void); /* never returns */ static void appTeardown(void); static void appSetup(int argc, char **argv); /* never returns */ static int appOptionsHandler(clientData cData, int opt_index, char *opt_arg); /* LOCAL VARIABLES */ /* index of first option that is not handled by the options handler. * If this is equal to argc, then input is from stdin. */ static int arg_index = 0; /* where to write the resulting set */ static skstream_t *set_stream = NULL; /* the compression method to use when writing the file. * sksiteCompmethodOptionsRegister() will set this to the default or * to the value the user specifies. */ static sk_compmethod_t comp_method; /* Holds the appOptionsEnum value for the operation being executed. */ static int operation = -1; /* Holds the appOptionsEnum value indicating whether --sample should * select based on a target sample size, or a ratio of the set size */ static int sample_type = 0; /* The random sample size from the --size switch */ static uint32_t sample_size; /* The random sample ratio from the --ratio switch */ static double sample_ratio; /* The PRNG seed */ static uint32_t sample_seed; /* Name of the output file */ static char *output_path = "stdout"; /* OPTIONS SETUP */ typedef enum { OPT_UNION, OPT_INTERSECT, OPT_DIFFERENCE, OPT_SAMPLE, OPT_SAMPLE_SIZE, OPT_SAMPLE_RATIO, OPT_SAMPLE_SEED, OPT_OUTPUT_PATH } appOptionsEnum; static struct option appOptions[] = { {"union", NO_ARG, 0, OPT_UNION}, {"intersect", NO_ARG, 0, OPT_INTERSECT}, {"difference", NO_ARG, 0, OPT_DIFFERENCE}, {"sample", NO_ARG, 0, OPT_SAMPLE}, {"size", REQUIRED_ARG, 0, OPT_SAMPLE_SIZE}, {"ratio", REQUIRED_ARG, 0, OPT_SAMPLE_RATIO}, {"seed", REQUIRED_ARG, 0, OPT_SAMPLE_SEED}, {"output-path", REQUIRED_ARG, 0, OPT_OUTPUT_PATH}, {0, 0, 0, 0} /* sentinel entry */ }; static const char *appHelp[] = { ("Create an IPset containing the IPs that exist in ANY of\n" "\tthe input IPsets"), ("Create an IPset containing the IPs that exist in ALL of\n" "\tthe input IPsets"), ("Create an IPset containing the IPs from the first IPset\n" "\tthat do not exist any subsequent IPset"), ("Create an IPset containing a random sample of IPs from\n" "\tall input IPsets. Requires the --size or --ratio switch."), ("Specify the sample size (number of IPs sampled from each\n" "\tinput IPset) for the --sample operation"), ("Specify the probability, as a floating point value between\n" "\t0.0 and 1.0, that an IP will be sampled"), "Specify the random number seed for the --sample operation", "Write the resulting IPset to this location. Def. stdout", (char *) NULL }; /* FUNCTION DEFINITIONS */ /* * appUsageLong(); * * Print complete usage information to USAGE_FH. Pass this * function to skOptionsSetUsageCallback(); optionsParse() will * call this funciton and then exit the program when the --help * option is given. */ static void appUsageLong(void) { #define USAGE_MSG \ ("{--union|--intersect|--difference|--sample} [SWITCHES] IPSET ...\n" \ "\tPerforms the specified operation on the input IPset file(s)\n" \ "\tand generates a new IPset file.\n") FILE *fh = USAGE_FH; skAppStandardUsage(fh, USAGE_MSG, appOptions, appHelp); sksiteCompmethodOptionsUsage(fh); } /* * appTeardown() * * Teardown all modules, close all files, and tidy up all * application state. * * This function is idempotent. */ static void appTeardown(void) { static int teardownFlag = 0; if (teardownFlag) { return; } teardownFlag = 1; if (set_stream) { skStreamDestroy(&set_stream); } skAppUnregister(); } /* * appSetup(argc, argv); * * Perform all the setup for this application include setting up * required modules, parsing options, etc. This function should be * passed the same arguments that were passed into main(). * * Returns to the caller if all setup succeeds. If anything fails, * this function will cause the application to exit with a FAILURE * exit status. */ static void appSetup(int argc, char **argv) { int rv; /* verify same number of options and help strings */ assert((sizeof(appHelp) / sizeof(char *)) == (sizeof(appOptions) / sizeof(struct option))); /* register the application */ skAppRegister(argv[0]); skOptionsSetUsageCallback(&appUsageLong); /* initialize variables */ set_stream = NULL; /* register the options */ if (optionsRegister(appOptions, (optHandler) appOptionsHandler, NULL) || sksiteCompmethodOptionsRegister(&comp_method)) { skAppPrintErr("unable to register options"); exit(EXIT_FAILURE); } /* parse options */ arg_index = optionsParse(argc, argv); assert(arg_index <= argc); if (arg_index < 0) { /* optionsParse() printed the error */ skAppUsage(); /* never returns */ } /* verify that we have something to do */ if (operation == -1) { skAppPrintErr("One of --%s, --%s, --%s, or --%s is required", appOptions[OPT_UNION].name, appOptions[OPT_INTERSECT].name, appOptions[OPT_DIFFERENCE].name, appOptions[OPT_SAMPLE].name); skAppUsage(); } /* either need name of set file(s) after options or a set file on stdin */ if ((arg_index == argc) && (FILEIsATty(stdin))) { skAppPrintErr("No files on the command line and" " stdin is connected to a terminal"); skAppUsage(); } /* verify that we have a sample size or ratio and seed the random * number generator */ if (operation == OPT_SAMPLE) { if (!sample_type) { skAppPrintErr(("The --%s switch requires a valid --%s or --%s " "argument"), appOptions[OPT_SAMPLE].name, appOptions[OPT_SAMPLE_SIZE].name, appOptions[OPT_SAMPLE_RATIO].name); skAppUsage(); } if (!sample_seed) { struct timeval tv; gettimeofday(&tv, NULL); sample_seed = (uint32_t)tv.tv_sec + (uint32_t)tv.tv_usec; } srandom((unsigned long)sample_seed); } /* open the output stream */ if ((rv = skStreamCreate(&set_stream, SK_IO_WRITE, SK_CONTENT_SILK)) || (rv = skStreamBind(set_stream, output_path)) || (rv = skStreamSetCompressionMethod(set_stream, comp_method)) || (rv = skStreamOpen(set_stream))) { skStreamPrintLastErr(set_stream, rv, &skAppPrintErr); skStreamDestroy(&set_stream); exit(EXIT_FAILURE); } if (atexit(appTeardown) < 0) { skAppPrintErr("unable to register appTeardown() with atexit()"); appTeardown(); exit(EXIT_FAILURE); } return; /* OK */ } /* * status = appOptionsHandler(cData, opt_index, opt_arg); * * This function is passed to optionsRegister(); it will be called * by optionsParse() for each user-specified switch that the * application has registered; it should handle the switch as * required---typically by setting global variables---and return 1 * if the switch processing failed or 0 if it succeeded. Returning * a non-zero from from the handler causes optionsParse() to return * a negative value. * * The clientData in 'cData' is typically ignored; 'opt_index' is * the index number that was specified as the last value for each * struct option in appOptions[]; 'opt_arg' is the user's argument * to the switch for options that have a REQUIRED_ARG or an * OPTIONAL_ARG. */ static int appOptionsHandler( clientData UNUSED(cData), int opt_index, char *opt_arg) { int rv; switch ((appOptionsEnum) opt_index) { case OPT_UNION: case OPT_INTERSECT: case OPT_DIFFERENCE: case OPT_SAMPLE: if ((operation >= 0) && (operation != opt_index)) { skAppPrintErr("Switches --%s and --%s are incompatible", appOptions[operation].name, appOptions[opt_index].name); return 1; } operation = opt_index; break; case OPT_SAMPLE_SIZE: if (sample_type) { skAppPrintErr("Switches --%s and --%s are incompatible", appOptions[sample_type].name, appOptions[opt_index].name); return 1; } sample_type = opt_index; if (skStringParseUint32(&sample_size, opt_arg, 1, 0)) { skAppPrintErr("Error parsing %s '%s'", appOptions[opt_index].name, opt_arg); return 1; } break; case OPT_SAMPLE_RATIO: if (sample_type) { skAppPrintErr("Switches %s and %s are incompatible", appOptions[sample_type].name, appOptions[opt_index].name); return 1; } sample_type = opt_index; rv = skStringParseDouble(&sample_ratio, opt_arg, 0.0, 1.0); if (rv) { if (rv < -10) { skAppPrintErr("The %s value '%s' is not between 0.0 and 1.0", appOptions[opt_index].name, opt_arg); } else { skAppPrintErr("Error parsing %s: '%s'", appOptions[opt_index].name, opt_arg); } return 1; } break; case OPT_SAMPLE_SEED: if (skStringParseUint32(&sample_seed, opt_arg, 1, 0)) { skAppPrintErr("Error parsing %s '%s'", appOptions[opt_index].name, opt_arg); skAppUsage(); } break; case OPT_OUTPUT_PATH: output_path = opt_arg; break; } return 0; /* OK */ } /* * filename = appNextInput(argc, argv); * * Return the name of the next input file from the command line or * the standard input if no files were given on the command line. */ static const char *appNextInput(int argc, char **argv) { static int initialized = 0; const char *fname = NULL; if (arg_index < argc) { /* get current file and prepare to get next */ fname = argv[arg_index]; ++arg_index; } else { if (initialized) { /* no more input */ return NULL; } /* input is from stdin */ fname = "stdin"; } initialized = 1; return fname; } static skIPTree_t *sampleSets(int argc, char **argv) { const char *filename = NULL; skIPTree_t *new_set = NULL; skIPTree_t *tmp_set = NULL; skIPTreeIterator_t iter; uint64_t set_count; uint64_t m = 0; uint64_t t = 1; uint32_t ipaddr; long frac; int rv; /* create the output set */ rv = skIPTreeCreate(&new_set); if (rv) { skAppPrintErr("Cannot create IPset: %s", skIPTreeStrError(rv)); return NULL; } while ((filename = appNextInput(argc, argv)) != NULL) { /* Open file and read IPset */ rv = skIPTreeLoad(filename, &tmp_set); if (rv) { skAppPrintErr("Error reading IPset from file '%s': %s", filename, skIPTreeStrError(rv)); skIPTreeDelete(&new_set); return NULL; } skIPTreeIteratorBind(&iter, tmp_set); /* TODO: Try list-based sampling algorithm */ if (sample_type == OPT_SAMPLE_SIZE) { set_count = skIPTreeCountIPs(tmp_set); m = t = 0; if (sample_size >= set_count) { /* add them all */ skIPTreeUnion(new_set, tmp_set); } else { while (skIPTreeIteratorNext(&ipaddr, &iter) == SK_ITERATOR_OK){ frac = (long)((double)RAND_MAX * (double)(sample_size - m) / (double)(set_count - (t - 1))); if (random() < frac) { skIPTreeAddAddress(ipaddr, new_set); ++m; if (m == sample_size) { break; } } ++t; } } } else { frac = (long)((double)RAND_MAX * sample_ratio); while (skIPTreeIteratorNext(&ipaddr, &iter) == SK_ITERATOR_OK) { if (random() < frac) { skIPTreeAddAddress(ipaddr, new_set); } } } skIPTreeDelete(&tmp_set); tmp_set = NULL; } return new_set; } int main(int argc, char **argv) { const char *filename = NULL; skIPTree_t *new_set = NULL; skIPTree_t *tmp_set = NULL; int rv; appSetup(argc, argv); /* never returns on error */ if (operation == OPT_SAMPLE) { new_set = sampleSets(argc, argv); if (new_set == NULL) { return 1; } } else { /* load the first set */ filename = appNextInput(argc, argv); rv = skIPTreeLoad(filename, &new_set); if (rv) { skAppPrintErr("Error reading IPset from file '%s': %s", filename, skIPTreeStrError(rv)); return 1; } while ((filename = appNextInput(argc, argv)) != NULL) { /* Open file and read IPset */ rv = skIPTreeLoad(filename, &tmp_set); if (rv) { skAppPrintErr("Error reading IPset from file '%s': %s", filename, skIPTreeStrError(rv)); return 1; } switch (operation) { case OPT_UNION: /* Do the union */ skIPTreeUnion(new_set, tmp_set); break; case OPT_INTERSECT: skIPTreeIntersect(new_set, tmp_set); break; case OPT_DIFFERENCE: skIPTreeSubtract(new_set, tmp_set); break; default: skAppPrintErr("Unhandled operation: '%d'" , operation); assert(0); exit(EXIT_FAILURE); break; } skIPTreeDelete(&tmp_set); tmp_set = NULL; } } /* write and close the set */ rv = skIPTreeWrite(set_stream, new_set); if (rv) { skAppPrintErr("Error writing IPset to file '%s': %s", skStreamGetPathname(set_stream), skIPTreeStrError(rv)); skStreamDestroy(&set_stream); skIPTreeDelete(&new_set); return 1; } if ((rv = skStreamClose(set_stream)) || (rv = skStreamDestroy(&set_stream))) { skStreamPrintLastErr(set_stream, rv, &skAppPrintErr); skStreamDestroy(&set_stream); skIPTreeDelete(&new_set); return 1; } skStreamDestroy(&set_stream); skIPTreeDelete(&new_set); /* done */ appTeardown(); return (0); } /* ** Local Variables: ** mode:c ** indent-tabs-mode:nil ** c-basic-offset:4 ** End: */