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