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