script/traffic_replay

   1 #!/usr/bin/env python3
   2 # Generates samba network traffic
   3 #
   4 # Copyright (C) Catalyst IT Ltd. 2017
   5 #
   6 # This program is free software; you can redistribute it and/or modify
   7 # it under the terms of the GNU General Public License as published by
   8 # the Free Software Foundation; either version 3 of the License, or
   9 # (at your option) any later version.
  10 #
  11 # This program is distributed in the hope that it will be useful,
  12 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  13 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  14 # GNU General Public License for more details.
  15 #
  16 # You should have received a copy of the GNU General Public License
  17 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
  18 #
  19 import sys
  20 import os
  21 import optparse
  22 import tempfile
  23 import shutil
  24 import random
  25
  26 sys.path.insert(0, "bin/python")
  27
  28 from samba import gensec, get_debug_level
  29 from samba.emulate import traffic
  30 import samba.getopt as options
  31 from samba.logger import get_samba_logger
  32 from samba.samdb import SamDB
  33 from samba.auth import system_session
  34
  35
  36 def print_err(*args, **kwargs):
  37     print(*args, file=sys.stderr, **kwargs)
  38
  39
  40 def main():
  41
  42     desc = ("Generates network traffic 'conversations' based on a model generated"
  43             " by script/traffic_learner. This traffic is sent to <dns-hostname>,"
  44             " which is the full DNS hostname of the DC being tested.")
  45
  46     parser = optparse.OptionParser(
  47         "%prog [--help|options] <model-file> <dns-hostname>",
  48         description=desc)
  49
  50     parser.add_option('--dns-rate', type='float', default=0,
  51                       help='fire extra DNS packets at this rate')
  52     parser.add_option('--dns-query-file', dest="dns_query_file",
  53                       help='A file contains DNS query list')
  54     parser.add_option('-B', '--badpassword-frequency',
  55                       type='float', default=0.0,
  56                       help='frequency of connections with bad passwords')
  57     parser.add_option('-K', '--prefer-kerberos',
  58                       action="store_true",
  59                       help='prefer kerberos when authenticating test users')
  60     parser.add_option('-I', '--instance-id', type='int', default=0,
  61                       help='Instance number, when running multiple instances')
  62     parser.add_option('-t', '--timing-data',
  63                       help=('write individual message timing data here '
  64                             '(- for stdout)'))
  65     parser.add_option('--preserve-tempdir', default=False, action="store_true",
  66                       help='do not delete temporary files')
  67     parser.add_option('-F', '--fixed-password',
  68                       type='string', default=None,
  69                       help=('Password used for the test users created. '
  70                             'Required'))
  71     parser.add_option('-c', '--clean-up',
  72                       action="store_true",
  73                       help='Clean up the generated groups and user accounts')
  74     parser.add_option('--random-seed', type='int', default=None,
  75                       help='Use to keep randomness consistent across multiple runs')
  76     parser.add_option('--stop-on-any-error',
  77                       action="store_true",
  78                       help='abort the whole thing if a child fails')
  79     model_group = optparse.OptionGroup(parser, 'Traffic Model Options',
  80                                        'These options alter the traffic '
  81                                        'generated by the model')
  82     model_group.add_option('-S', '--scale-traffic', type='float',
  83                            help=('Increase the number of conversations by '
  84                                  'this factor (or use -T)'))
  85     parser.add_option('-T', '--packets-per-second', type=float,
  86                       help=('attempt this many packets per second '
  87                             '(alternative to -S)'))
  88     parser.add_option('--old-scale',
  89                       action="store_true",
  90                       help='emulate the old scale for traffic')
  91     model_group.add_option('-D', '--duration', type='float', default=60.0,
  92                            help=('Run model for this long (approx). '
  93                                  'Default 60s for models'))
  94     model_group.add_option('--latency-timeout', type='float', default=None,
  95                            help=('Wait this long for last packet to finish'))
  96     model_group.add_option('-r', '--replay-rate', type='float', default=1.0,
  97                            help='Replay the traffic faster by this factor')
  98     model_group.add_option('--conversation-persistence', type='float',
  99                            default=0.0,
 100                            help=('chance (0 to 1) that a conversation waits '
 101                                  'when it would have died'))
 102     model_group.add_option('--traffic-summary',
 103                            help=('Generate a traffic summary file and write '
 104                                  'it here (- for stdout)'))
 105     parser.add_option_group(model_group)
 106
 107     user_gen_group = optparse.OptionGroup(parser, 'Generate User Options',
 108                                           "Add extra user/groups on the DC to "
 109                                           "increase the DB size. These extra "
 110                                           "users aren't used for traffic "
 111                                           "generation.")
 112     user_gen_group.add_option('-G', '--generate-users-only',
 113                               action="store_true",
 114                               help='Generate the users, but do not replay '
 115                               'the traffic')
 116     user_gen_group.add_option('-n', '--number-of-users', type='int', default=0,
 117                               help='Total number of test users to create')
 118     user_gen_group.add_option('--number-of-groups', type='int', default=None,
 119                               help='Create this many groups')
 120     user_gen_group.add_option('--average-groups-per-user',
 121                               type='int', default=0,
 122                               help='Assign the test users to this '
 123                               'many groups on average')
 124     user_gen_group.add_option('--group-memberships', type='int', default=0,
 125                               help='Total memberships to assign across all '
 126                               'test users and all groups')
 127     user_gen_group.add_option('--max-members', type='int', default=None,
 128                               help='Max users to add to any one group')
 129     parser.add_option_group(user_gen_group)
 130
 131     sambaopts = options.SambaOptions(parser)
 132     parser.add_option_group(sambaopts)
 133     parser.add_option_group(options.VersionOptions(parser))
 134     credopts = options.CredentialsOptions(parser)
 135     parser.add_option_group(credopts)
 136
 137     # the --no-password credential doesn't make sense for this tool
 138     if parser.has_option('-N'):
 139         parser.remove_option('-N')
 140
 141     opts, args = parser.parse_args()
 142
 143     # First ensure we have reasonable arguments
 144
 145     if len(args) == 1:
 146         model_file = None
 147         host    = args[0]
 148     elif len(args) == 2:
 149         model_file, host = args
 150     else:
 151         parser.print_usage()
 152         return
 153
 154     lp = sambaopts.get_loadparm()
 155     debuglevel = get_debug_level()
 156     logger = get_samba_logger(name=__name__,
 157                               verbose=debuglevel > 3,
 158                               quiet=debuglevel < 1)
 159
 160     traffic.DEBUG_LEVEL = debuglevel
 161     # pass log level down to traffic module to make sure level is controlled
 162     traffic.LOGGER.setLevel(logger.getEffectiveLevel())
 163
 164     if opts.clean_up:
 165         logger.info("Removing user and machine accounts")
 166         lp    = sambaopts.get_loadparm()
 167         creds = credopts.get_credentials(lp)
 168         creds.set_gensec_features(creds.get_gensec_features() | gensec.FEATURE_SEAL)
 169         ldb   = traffic.openLdb(host, creds, lp)
 170         traffic.clean_up_accounts(ldb, opts.instance_id)
 171         exit(0)
 172
 173     if model_file:
 174         if not os.path.exists(model_file):
 175             logger.error("Model file %s doesn't exist" % model_file)
 176             sys.exit(1)
 177     # the model-file can be ommitted for --generate-users-only and
 178     # --cleanup-up, but it should be specified in all other cases
 179     elif not opts.generate_users_only:
 180         logger.error("No model file specified to replay traffic from")
 181         sys.exit(1)
 182
 183     if not opts.fixed_password:
 184         logger.error(("Please use --fixed-password to specify a password"
 185                       " for the users created as part of this test"))
 186         sys.exit(1)
 187
 188     if opts.random_seed is not None:
 189         random.seed(opts.random_seed)
 190
 191     creds = credopts.get_credentials(lp)
 192     creds.set_gensec_features(creds.get_gensec_features() | gensec.FEATURE_SEAL)
 193
 194     domain = creds.get_domain()
 195     if domain:
 196         lp.set("workgroup", domain)
 197     else:
 198         domain = lp.get("workgroup")
 199         if domain == "WORKGROUP":
 200             logger.error(("NETBIOS domain does not appear to be "
 201                           "specified, use the --workgroup option"))
 202             sys.exit(1)
 203
 204     if not opts.realm and not lp.get('realm'):
 205         logger.error("Realm not specified, use the --realm option")
 206         sys.exit(1)
 207
 208     if opts.generate_users_only and not (opts.number_of_users or
 209                                          opts.number_of_groups):
 210         logger.error(("Please specify the number of users and/or groups "
 211                       "to generate."))
 212         sys.exit(1)
 213
 214     if opts.group_memberships and opts.average_groups_per_user:
 215         logger.error(("--group-memberships and --average-groups-per-user"
 216                       " are incompatible options - use one or the other"))
 217         sys.exit(1)
 218
 219     if not opts.number_of_groups and opts.average_groups_per_user:
 220         logger.error(("--average-groups-per-user requires "
 221                       "--number-of-groups"))
 222         sys.exit(1)
 223
 224     if opts.number_of_groups and opts.average_groups_per_user:
 225         if opts.number_of_groups < opts.average_groups_per_user:
 226             logger.error(("--average-groups-per-user can not be more than "
 227                           "--number-of-groups"))
 228             sys.exit(1)
 229
 230     if not opts.number_of_groups and opts.group_memberships:
 231         logger.error("--group-memberships requires --number-of-groups")
 232         sys.exit(1)
 233
 234     if opts.scale_traffic is not None and opts.packets_per_second is not None:
 235         logger.error("--scale-traffic and --packets-per-second "
 236                      "are incompatible. Use one or the other.")
 237         sys.exit(1)
 238
 239     if not opts.scale_traffic and not opts.packets_per_second:
 240         logger.info("No packet rate specified. Using --scale-traffic=1.0")
 241         opts.scale_traffic = 1.0
 242
 243     if opts.timing_data not in ('-', None):
 244         try:
 245             open(opts.timing_data, 'w').close()
 246         except IOError:
 247             # exception info will be added to log automatically
 248             logger.exception(("the supplied timing data destination "
 249                               "(%s) is not writable" % opts.timing_data))
 250             sys.exit()
 251
 252     if opts.traffic_summary not in ('-', None):
 253         try:
 254             open(opts.traffic_summary, 'w').close()
 255         except IOError:
 256             # exception info will be added to log automatically
 257             if debuglevel > 0:
 258                 import traceback
 259                 traceback.print_exc()
 260             logger.exception(("the supplied traffic summary destination "
 261                               "(%s) is not writable" % opts.traffic_summary))
 262             sys.exit()
 263
 264     if opts.old_scale:
 265         # we used to use a silly calculation based on the number
 266         # of conversations; now we use the number of packets and
 267         # scale traffic accurately. To roughly compare with older
 268         # numbers you use --old-scale which approximates as follows:
 269         opts.scale_traffic *= 0.55
 270
 271     # ingest the model
 272     if model_file and not opts.generate_users_only:
 273         model = traffic.TrafficModel()
 274         try:
 275             model.load(model_file)
 276         except ValueError:
 277             if debuglevel > 0:
 278                 import traceback
 279                 traceback.print_exc()
 280             logger.error(("Could not parse %s, which does not seem to be "
 281                           "a model generated by script/traffic_learner."
 282                           % model_file))
 283             sys.exit(1)
 284
 285         logger.info(("Using the specified model file to "
 286                      "generate conversations"))
 287
 288         if opts.scale_traffic:
 289             packets_per_second = model.scale_to_packet_rate(opts.scale_traffic)
 290         else:
 291             packets_per_second =  opts.packets_per_second
 292
 293         conversations = \
 294             model.generate_conversation_sequences(
 295                 packets_per_second,
 296                 opts.duration,
 297                 opts.replay_rate,
 298                 opts.conversation_persistence)
 299     else:
 300         conversations = []
 301
 302     if opts.number_of_users and opts.number_of_users < len(conversations):
 303         logger.error(("--number-of-users (%d) is less than the "
 304                       "number of conversations to replay (%d)"
 305                      % (opts.number_of_users, len(conversations))))
 306         sys.exit(1)
 307
 308     number_of_users = max(opts.number_of_users, len(conversations))
 309
 310     if opts.number_of_groups is None:
 311         opts.number_of_groups = max(int(number_of_users / 10), 1)
 312
 313     max_memberships = number_of_users * opts.number_of_groups
 314
 315     if not opts.group_memberships and opts.average_groups_per_user:
 316         opts.group_memberships = opts.average_groups_per_user * number_of_users
 317         logger.info(("Using %d group-memberships based on %u average "
 318                      "memberships for %d users"
 319                      % (opts.group_memberships,
 320                         opts.average_groups_per_user, number_of_users)))
 321
 322     if opts.group_memberships > max_memberships:
 323         logger.error(("The group memberships specified (%d) exceeds "
 324                       "the total users (%d) * total groups (%d)"
 325                       % (opts.group_memberships, number_of_users,
 326                          opts.number_of_groups)))
 327         sys.exit(1)
 328
 329     # if no groups were specified by the user, then make sure we create some
 330     # group memberships (otherwise it's not really a fair test)
 331     if not opts.group_memberships and not opts.average_groups_per_user:
 332         opts.group_memberships = min(number_of_users * 5, max_memberships)
 333
 334     # Get an LDB connection.
 335     try:
 336         # if we're only adding users, then it's OK to pass a sam.ldb filepath
 337         # as the host, which creates the users much faster. In all other cases
 338         # we should be connecting to a remote DC
 339         if opts.generate_users_only and os.path.isfile(host):
 340             ldb = SamDB(url="ldb://{0}".format(host),
 341                         session_info=system_session(), lp=lp)
 342         else:
 343             ldb = traffic.openLdb(host, creds, lp)
 344     except:
 345         logger.error(("\nInitial LDAP connection failed! Did you supply "
 346                       "a DNS host name and the correct credentials?"))
 347         sys.exit(1)
 348
 349     if opts.generate_users_only:
 350         # generate computer accounts for added realism. Assume there will be
 351         # some overhang with more computer accounts than users
 352         computer_accounts = int(1.25 * number_of_users)
 353         traffic.generate_users_and_groups(ldb,
 354                                           opts.instance_id,
 355                                           opts.fixed_password,
 356                                           opts.number_of_users,
 357                                           opts.number_of_groups,
 358                                           opts.group_memberships,
 359                                           opts.max_members,
 360                                           machine_accounts=computer_accounts,
 361                                           traffic_accounts=False)
 362         sys.exit()
 363
 364     tempdir = tempfile.mkdtemp(prefix="samba_tg_")
 365     logger.info("Using temp dir %s" % tempdir)
 366
 367     traffic.generate_users_and_groups(ldb,
 368                                       opts.instance_id,
 369                                       opts.fixed_password,
 370                                       number_of_users,
 371                                       opts.number_of_groups,
 372                                       opts.group_memberships,
 373                                       opts.max_members,
 374                                       machine_accounts=len(conversations),
 375                                       traffic_accounts=True)
 376
 377     accounts = traffic.generate_replay_accounts(ldb,
 378                                                 opts.instance_id,
 379                                                 len(conversations),
 380                                                 opts.fixed_password)
 381
 382     statsdir = traffic.mk_masked_dir(tempdir, 'stats')
 383
 384     if opts.traffic_summary:
 385         if opts.traffic_summary == '-':
 386             summary_dest = sys.stdout
 387         else:
 388             summary_dest = open(opts.traffic_summary, 'w')
 389
 390         logger.info("Writing traffic summary")
 391         summaries = []
 392         for c in traffic.seq_to_conversations(conversations):
 393             summaries += c.replay_as_summary_lines()
 394
 395         summaries.sort()
 396         for (time, line) in summaries:
 397             print(line, file=summary_dest)
 398
 399         exit(0)
 400
 401     traffic.replay(conversations,
 402                    host,
 403                    lp=lp,
 404                    creds=creds,
 405                    accounts=accounts,
 406                    dns_rate=opts.dns_rate,
 407                    dns_query_file=opts.dns_query_file,
 408                    duration=opts.duration,
 409                    latency_timeout=opts.latency_timeout,
 410                    badpassword_frequency=opts.badpassword_frequency,
 411                    prefer_kerberos=opts.prefer_kerberos,
 412                    statsdir=statsdir,
 413                    domain=domain,
 414                    base_dn=ldb.domain_dn(),
 415                    ou=traffic.ou_name(ldb, opts.instance_id),
 416                    tempdir=tempdir,
 417                    stop_on_any_error=opts.stop_on_any_error,
 418                    domain_sid=ldb.get_domain_sid(),
 419                    instance_id=opts.instance_id)
 420
 421     if opts.timing_data == '-':
 422         timing_dest = sys.stdout
 423     elif opts.timing_data is None:
 424         timing_dest = None
 425     else:
 426         timing_dest = open(opts.timing_data, 'w')
 427
 428     logger.info("Generating statistics")
 429     traffic.generate_stats(statsdir, timing_dest)
 430
 431     if not opts.preserve_tempdir:
 432         logger.info("Removing temporary directory")
 433         shutil.rmtree(tempdir)
 434     else:
 435         # delete the empty directories anyway. There are thousands of
 436         # them and they're EMPTY.
 437         for d in os.listdir(tempdir):
 438             if d.startswith('conversation-'):
 439                 path = os.path.join(tempdir, d)
 440                 try:
 441                     os.rmdir(path)
 442                 except OSError as e:
 443                     logger.info("not removing %s (%s)" % (path, e))
 444
 445 main()