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