script/traffic_replay

   1 #!/usr/bin/env python
   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 from __future__ import print_function
  20 import sys
  21 import os
  22 import optparse
  23 import tempfile
  24 import shutil
  25 import random
  26
  27 sys.path.insert(0, "bin/python")
  28
  29 from samba import gensec, get_debug_level
  30 from samba.emulate import traffic
  31 import samba.getopt as options
  32 from samba.logger import get_samba_logger
  33 from samba.samdb import SamDB
  34 from samba.auth import system_session
  35
  36
  37 def print_err(*args, **kwargs):
  38     print(*args, file=sys.stderr, **kwargs)
  39
  40
  41 def main():
  42
  43     desc = ("Generates network traffic 'conversations' based on <summary-file>"
  44             " (which should be the output file produced by either traffic_learner"
  45             " or traffic_summary.pl). This traffic is sent to <dns-hostname>,"
  46             " which is the full DNS hostname of the DC being tested.")
  47
  48     parser = optparse.OptionParser(
  49         "%prog [--help|options] <summary-file> <dns-hostname>",
  50         description=desc)
  51
  52     parser.add_option('--dns-rate', type='float', default=0,
  53                       help='fire extra DNS packets at this rate')
  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=0,
  75                       help='Use to keep randomness consistent across multiple runs')
  76
  77     model_group = optparse.OptionGroup(parser, 'Traffic Model Options',
  78                                        'These options alter the traffic '
  79                                        'generated when the summary-file is a '
  80                                        'traffic-model (produced by '
  81                                        'traffic_learner)')
  82     model_group.add_option('-S', '--scale-traffic', type='float', default=1.0,
  83                            help='Increase the number of conversations by '
  84                            'this factor')
  85     model_group.add_option('-D', '--duration', type='float', default=None,
  86                            help=('Run model for this long (approx). '
  87                                  'Default 60s for models'))
  88     model_group.add_option('-r', '--replay-rate', type='float', default=1.0,
  89                            help='Replay the traffic faster by this factor')
  90     model_group.add_option('--traffic-summary',
  91                            help=('Generate a traffic summary file and write '
  92                                  'it here (- for stdout)'))
  93     parser.add_option_group(model_group)
  94
  95     user_gen_group = optparse.OptionGroup(parser, 'Generate User Options',
  96                                           "Add extra user/groups on the DC to "
  97                                           "increase the DB size. These extra "
  98                                           "users aren't used for traffic "
  99                                           "generation.")
 100     user_gen_group.add_option('-G', '--generate-users-only',
 101                               action="store_true",
 102                               help='Generate the users, but do not replay '
 103                               'the traffic')
 104     user_gen_group.add_option('-n', '--number-of-users', type='int', default=0,
 105                               help='Total number of test users to create')
 106     user_gen_group.add_option('--number-of-groups', type='int', default=0,
 107                               help='Create this many groups')
 108     user_gen_group.add_option('--average-groups-per-user',
 109                               type='int', default=0,
 110                               help='Assign the test users to this '
 111                               'many groups on average')
 112     user_gen_group.add_option('--group-memberships', type='int', default=0,
 113                               help='Total memberships to assign across all '
 114                               'test users and all groups')
 115     parser.add_option_group(user_gen_group)
 116
 117     sambaopts = options.SambaOptions(parser)
 118     parser.add_option_group(sambaopts)
 119     parser.add_option_group(options.VersionOptions(parser))
 120     credopts = options.CredentialsOptions(parser)
 121     parser.add_option_group(credopts)
 122
 123     # the --no-password credential doesn't make sense for this tool
 124     if parser.has_option('-N'):
 125         parser.remove_option('-N')
 126
 127     opts, args = parser.parse_args()
 128
 129     # First ensure we have reasonable arguments
 130
 131     if len(args) == 1:
 132         summary = None
 133         host    = args[0]
 134     elif len(args) == 2:
 135         summary, host = args
 136     else:
 137         parser.print_usage()
 138         return
 139
 140     lp = sambaopts.get_loadparm()
 141     debuglevel = get_debug_level()
 142     logger = get_samba_logger(name=__name__,
 143                               verbose=debuglevel > 3,
 144                               quiet=debuglevel < 1)
 145
 146     traffic.DEBUG_LEVEL = debuglevel
 147     # pass log level down to traffic module to make sure level is controlled
 148     traffic.LOGGER.setLevel(logger.getEffectiveLevel())
 149
 150     if opts.clean_up:
 151         logger.info("Removing user and machine accounts")
 152         lp    = sambaopts.get_loadparm()
 153         creds = credopts.get_credentials(lp)
 154         creds.set_gensec_features(creds.get_gensec_features() | gensec.FEATURE_SEAL)
 155         ldb   = traffic.openLdb(host, creds, lp)
 156         traffic.clean_up_accounts(ldb, opts.instance_id)
 157         exit(0)
 158
 159     if summary:
 160         if not os.path.exists(summary):
 161             logger.error("Summary file %s doesn't exist" % summary)
 162             sys.exit(1)
 163     # the summary-file can be ommitted for --generate-users-only and
 164     # --cleanup-up, but it should be specified in all other cases
 165     elif not opts.generate_users_only:
 166         logger.error("No summary-file specified to replay traffic from")
 167         sys.exit(1)
 168
 169     if not opts.fixed_password:
 170         logger.error(("Please use --fixed-password to specify a password"
 171                       " for the users created as part of this test"))
 172         sys.exit(1)
 173
 174     if opts.random_seed:
 175         random.seed(opts.random_seed)
 176
 177     creds = credopts.get_credentials(lp)
 178     creds.set_gensec_features(creds.get_gensec_features() | gensec.FEATURE_SEAL)
 179
 180     domain = creds.get_domain()
 181     if domain:
 182         lp.set("workgroup", domain)
 183     else:
 184         domain = lp.get("workgroup")
 185         if domain == "WORKGROUP":
 186             logger.error(("NETBIOS domain does not appear to be "
 187                           "specified, use the --workgroup option"))
 188             sys.exit(1)
 189
 190     if not opts.realm and not lp.get('realm'):
 191         logger.error("Realm not specified, use the --realm option")
 192         sys.exit(1)
 193
 194     if opts.generate_users_only and not (opts.number_of_users or
 195                                          opts.number_of_groups):
 196         logger.error(("Please specify the number of users and/or groups "
 197                       "to generate."))
 198         sys.exit(1)
 199
 200     if opts.group_memberships and opts.average_groups_per_user:
 201         logger.error(("--group-memberships and --average-groups-per-user"
 202                       " are incompatible options - use one or the other"))
 203         sys.exit(1)
 204
 205     if not opts.number_of_groups and opts.average_groups_per_user:
 206         logger.error(("--average-groups-per-user requires "
 207                       "--number-of-groups"))
 208         sys.exit(1)
 209
 210     if opts.number_of_groups and opts.average_groups_per_user:
 211         if opts.number_of_groups < opts.average_groups_per_user:
 212             logger.error(("--average-groups-per-user can not be more than "
 213                           "--number-of-groups"))
 214             sys.exit(1)
 215
 216     if not opts.number_of_groups and opts.group_memberships:
 217         logger.error("--group-memberships requires --number-of-groups")
 218         sys.exit(1)
 219
 220     if opts.timing_data not in ('-', None):
 221         try:
 222             open(opts.timing_data, 'w').close()
 223         except IOError:
 224             # exception info will be added to log automatically
 225             logger.exception(("the supplied timing data destination "
 226                               "(%s) is not writable" % opts.timing_data))
 227             sys.exit()
 228
 229     if opts.traffic_summary not in ('-', None):
 230         try:
 231             open(opts.traffic_summary, 'w').close()
 232         except IOError:
 233             # exception info will be added to log automatically
 234             logger.exception(("the supplied traffic summary destination "
 235                               "(%s) is not writable" % opts.traffic_summary))
 236             sys.exit()
 237
 238     duration = opts.duration
 239     if duration is None:
 240         duration = 60.0
 241
 242     # ingest the model or traffic summary
 243     if summary:
 244         try:
 245             conversations, interval, duration, dns_counts = \
 246                                             traffic.ingest_summaries([summary])
 247
 248             logger.info(("Using conversations from the traffic summary "
 249                          "file specified"))
 250
 251             # honour the specified duration if it's different to the
 252             # capture duration
 253             if opts.duration is not None:
 254                 duration = opts.duration
 255
 256         except ValueError as e:
 257             if not str(e).startswith('need more than'):
 258                 raise
 259
 260             model = traffic.TrafficModel()
 261
 262             try:
 263                 model.load(summary)
 264             except ValueError:
 265                 logger.error(("Could not parse %s. The summary file "
 266                               "should be the output from either the "
 267                               "traffic_summary.pl or "
 268                               "traffic_learner scripts.") % summary)
 269                 sys.exit()
 270
 271             logger.info(("Using the specified model file to "
 272                          "generate conversations"))
 273
 274             conversations = model.generate_conversations(opts.scale_traffic,
 275                                                          duration,
 276                                                          opts.replay_rate)
 277
 278     else:
 279         conversations = []
 280
 281     if debuglevel > 5:
 282         for c in conversations:
 283             for p in c.packets:
 284                 print("    ", p, file=sys.stderr)
 285
 286         print('=' * 72, file=sys.stderr)
 287
 288     if opts.number_of_users and opts.number_of_users < len(conversations):
 289         logger.error(("--number-of-users (%d) is less than the "
 290                       "number of conversations to replay (%d)"
 291                      % (opts.number_of_users, len(conversations))))
 292         sys.exit(1)
 293
 294     number_of_users = max(opts.number_of_users, len(conversations))
 295     max_memberships = number_of_users * opts.number_of_groups
 296
 297     if not opts.group_memberships and opts.average_groups_per_user:
 298         opts.group_memberships = opts.average_groups_per_user * number_of_users
 299         logger.info(("Using %d group-memberships based on %u average "
 300                      "memberships for %d users"
 301                      % (opts.group_memberships,
 302                         opts.average_groups_per_user, number_of_users)))
 303
 304     if opts.group_memberships > max_memberships:
 305         logger.error(("The group memberships specified (%d) exceeds "
 306                       "the total users (%d) * total groups (%d)"
 307                       % (opts.group_memberships, number_of_users,
 308                          opts.number_of_groups)))
 309         sys.exit(1)
 310
 311     # Get an LDB connection.
 312     try:
 313         # if we're only adding users, then it's OK to pass a sam.ldb filepath
 314         # as the host, which creates the users much faster. In all other cases
 315         # we should be connecting to a remote DC
 316         if opts.generate_users_only and os.path.isfile(host):
 317             ldb = SamDB(url="ldb://{0}".format(host),
 318                         session_info=system_session(), lp=lp)
 319         else:
 320             ldb = traffic.openLdb(host, creds, lp)
 321     except:
 322         logger.error(("\nInitial LDAP connection failed! Did you supply "
 323                       "a DNS host name and the correct credentials?"))
 324         sys.exit(1)
 325
 326     if opts.generate_users_only:
 327         traffic.generate_users_and_groups(ldb,
 328                                           opts.instance_id,
 329                                           opts.fixed_password,
 330                                           opts.number_of_users,
 331                                           opts.number_of_groups,
 332                                           opts.group_memberships)
 333         sys.exit()
 334
 335     tempdir = tempfile.mkdtemp(prefix="samba_tg_")
 336     logger.info("Using temp dir %s" % tempdir)
 337
 338     traffic.generate_users_and_groups(ldb,
 339                                       opts.instance_id,
 340                                       opts.fixed_password,
 341                                       number_of_users,
 342                                       opts.number_of_groups,
 343                                       opts.group_memberships)
 344
 345     accounts = traffic.generate_replay_accounts(ldb,
 346                                                 opts.instance_id,
 347                                                 len(conversations),
 348                                                 opts.fixed_password)
 349
 350     statsdir = traffic.mk_masked_dir(tempdir, 'stats')
 351
 352     if opts.traffic_summary:
 353         if opts.traffic_summary == '-':
 354             summary_dest = sys.stdout
 355         else:
 356             summary_dest = open(opts.traffic_summary, 'w')
 357
 358         logger.info("Writing traffic summary")
 359         summaries = []
 360         for c in conversations:
 361             summaries += c.replay_as_summary_lines()
 362
 363         summaries.sort()
 364         for (time, line) in summaries:
 365             print(line, file=summary_dest)
 366
 367         exit(0)
 368
 369     traffic.replay(conversations, host,
 370                    lp=lp,
 371                    creds=creds,
 372                    accounts=accounts,
 373                    dns_rate=opts.dns_rate,
 374                    duration=duration,
 375                    badpassword_frequency=opts.badpassword_frequency,
 376                    prefer_kerberos=opts.prefer_kerberos,
 377                    statsdir=statsdir,
 378                    domain=domain,
 379                    base_dn=ldb.domain_dn(),
 380                    ou=traffic.ou_name(ldb, opts.instance_id),
 381                    tempdir=tempdir,
 382                    domain_sid=ldb.get_domain_sid())
 383
 384     if opts.timing_data == '-':
 385         timing_dest = sys.stdout
 386     elif opts.timing_data is None:
 387         timing_dest = None
 388     else:
 389         timing_dest = open(opts.timing_data, 'w')
 390
 391     logger.info("Generating statistics")
 392     traffic.generate_stats(statsdir, timing_dest)
 393
 394     if not opts.preserve_tempdir:
 395         logger.info("Removing temporary directory")
 396         shutil.rmtree(tempdir)
 397
 398
 399 main()