PY3: change shebang to python3 in script dir
[garming/samba-autobuild/.git] / 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 <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     user_gen_group.add_option('--max-members', type='int', default=None,
116                               help='Max users to add to any one group')
117     parser.add_option_group(user_gen_group)
118
119     sambaopts = options.SambaOptions(parser)
120     parser.add_option_group(sambaopts)
121     parser.add_option_group(options.VersionOptions(parser))
122     credopts = options.CredentialsOptions(parser)
123     parser.add_option_group(credopts)
124
125     # the --no-password credential doesn't make sense for this tool
126     if parser.has_option('-N'):
127         parser.remove_option('-N')
128
129     opts, args = parser.parse_args()
130
131     # First ensure we have reasonable arguments
132
133     if len(args) == 1:
134         summary = None
135         host    = args[0]
136     elif len(args) == 2:
137         summary, host = args
138     else:
139         parser.print_usage()
140         return
141
142     lp = sambaopts.get_loadparm()
143     debuglevel = get_debug_level()
144     logger = get_samba_logger(name=__name__,
145                               verbose=debuglevel > 3,
146                               quiet=debuglevel < 1)
147
148     traffic.DEBUG_LEVEL = debuglevel
149     # pass log level down to traffic module to make sure level is controlled
150     traffic.LOGGER.setLevel(logger.getEffectiveLevel())
151
152     if opts.clean_up:
153         logger.info("Removing user and machine accounts")
154         lp    = sambaopts.get_loadparm()
155         creds = credopts.get_credentials(lp)
156         creds.set_gensec_features(creds.get_gensec_features() | gensec.FEATURE_SEAL)
157         ldb   = traffic.openLdb(host, creds, lp)
158         traffic.clean_up_accounts(ldb, opts.instance_id)
159         exit(0)
160
161     if summary:
162         if not os.path.exists(summary):
163             logger.error("Summary file %s doesn't exist" % summary)
164             sys.exit(1)
165     # the summary-file can be ommitted for --generate-users-only and
166     # --cleanup-up, but it should be specified in all other cases
167     elif not opts.generate_users_only:
168         logger.error("No summary-file specified to replay traffic from")
169         sys.exit(1)
170
171     if not opts.fixed_password:
172         logger.error(("Please use --fixed-password to specify a password"
173                       " for the users created as part of this test"))
174         sys.exit(1)
175
176     if opts.random_seed:
177         random.seed(opts.random_seed)
178
179     creds = credopts.get_credentials(lp)
180     creds.set_gensec_features(creds.get_gensec_features() | gensec.FEATURE_SEAL)
181
182     domain = creds.get_domain()
183     if domain:
184         lp.set("workgroup", domain)
185     else:
186         domain = lp.get("workgroup")
187         if domain == "WORKGROUP":
188             logger.error(("NETBIOS domain does not appear to be "
189                           "specified, use the --workgroup option"))
190             sys.exit(1)
191
192     if not opts.realm and not lp.get('realm'):
193         logger.error("Realm not specified, use the --realm option")
194         sys.exit(1)
195
196     if opts.generate_users_only and not (opts.number_of_users or
197                                          opts.number_of_groups):
198         logger.error(("Please specify the number of users and/or groups "
199                       "to generate."))
200         sys.exit(1)
201
202     if opts.group_memberships and opts.average_groups_per_user:
203         logger.error(("--group-memberships and --average-groups-per-user"
204                       " are incompatible options - use one or the other"))
205         sys.exit(1)
206
207     if not opts.number_of_groups and opts.average_groups_per_user:
208         logger.error(("--average-groups-per-user requires "
209                       "--number-of-groups"))
210         sys.exit(1)
211
212     if opts.number_of_groups and opts.average_groups_per_user:
213         if opts.number_of_groups < opts.average_groups_per_user:
214             logger.error(("--average-groups-per-user can not be more than "
215                           "--number-of-groups"))
216             sys.exit(1)
217
218     if not opts.number_of_groups and opts.group_memberships:
219         logger.error("--group-memberships requires --number-of-groups")
220         sys.exit(1)
221
222     if opts.timing_data not in ('-', None):
223         try:
224             open(opts.timing_data, 'w').close()
225         except IOError:
226             # exception info will be added to log automatically
227             logger.exception(("the supplied timing data destination "
228                               "(%s) is not writable" % opts.timing_data))
229             sys.exit()
230
231     if opts.traffic_summary not in ('-', None):
232         try:
233             open(opts.traffic_summary, 'w').close()
234         except IOError:
235             # exception info will be added to log automatically
236             logger.exception(("the supplied traffic summary destination "
237                               "(%s) is not writable" % opts.traffic_summary))
238             sys.exit()
239
240     duration = opts.duration
241     if duration is None:
242         duration = 60.0
243
244     # ingest the model or traffic summary
245     if summary:
246         try:
247             conversations, interval, duration, dns_counts = \
248                                             traffic.ingest_summaries([summary])
249
250             logger.info(("Using conversations from the traffic summary "
251                          "file specified"))
252
253             # honour the specified duration if it's different to the
254             # capture duration
255             if opts.duration is not None:
256                 duration = opts.duration
257
258         except ValueError as e:
259             if not str(e).startswith('need more than'):
260                 raise
261
262             model = traffic.TrafficModel()
263
264             try:
265                 model.load(summary)
266             except ValueError:
267                 logger.error(("Could not parse %s. The summary file "
268                               "should be the output from either the "
269                               "traffic_summary.pl or "
270                               "traffic_learner scripts.") % summary)
271                 sys.exit()
272
273             logger.info(("Using the specified model file to "
274                          "generate conversations"))
275
276             conversations = model.generate_conversations(opts.scale_traffic,
277                                                          duration,
278                                                          opts.replay_rate)
279
280     else:
281         conversations = []
282
283     if debuglevel > 5:
284         for c in conversations:
285             for p in c.packets:
286                 print("    ", p, file=sys.stderr)
287
288         print('=' * 72, file=sys.stderr)
289
290     if opts.number_of_users and opts.number_of_users < len(conversations):
291         logger.error(("--number-of-users (%d) is less than the "
292                       "number of conversations to replay (%d)"
293                      % (opts.number_of_users, len(conversations))))
294         sys.exit(1)
295
296     number_of_users = max(opts.number_of_users, len(conversations))
297     max_memberships = number_of_users * opts.number_of_groups
298
299     if not opts.group_memberships and opts.average_groups_per_user:
300         opts.group_memberships = opts.average_groups_per_user * number_of_users
301         logger.info(("Using %d group-memberships based on %u average "
302                      "memberships for %d users"
303                      % (opts.group_memberships,
304                         opts.average_groups_per_user, number_of_users)))
305
306     if opts.group_memberships > max_memberships:
307         logger.error(("The group memberships specified (%d) exceeds "
308                       "the total users (%d) * total groups (%d)"
309                       % (opts.group_memberships, number_of_users,
310                          opts.number_of_groups)))
311         sys.exit(1)
312
313     # Get an LDB connection.
314     try:
315         # if we're only adding users, then it's OK to pass a sam.ldb filepath
316         # as the host, which creates the users much faster. In all other cases
317         # we should be connecting to a remote DC
318         if opts.generate_users_only and os.path.isfile(host):
319             ldb = SamDB(url="ldb://{0}".format(host),
320                         session_info=system_session(), lp=lp)
321         else:
322             ldb = traffic.openLdb(host, creds, lp)
323     except:
324         logger.error(("\nInitial LDAP connection failed! Did you supply "
325                       "a DNS host name and the correct credentials?"))
326         sys.exit(1)
327
328     if opts.generate_users_only:
329         # generate computer accounts for added realism. Assume there will be
330         # some overhang with more computer accounts than users
331         computer_accounts = int(1.25 * number_of_users)
332         traffic.generate_users_and_groups(ldb,
333                                           opts.instance_id,
334                                           opts.fixed_password,
335                                           opts.number_of_users,
336                                           opts.number_of_groups,
337                                           opts.group_memberships,
338                                           opts.max_members,
339                                           machine_accounts=computer_accounts,
340                                           traffic_accounts=False)
341         sys.exit()
342
343     tempdir = tempfile.mkdtemp(prefix="samba_tg_")
344     logger.info("Using temp dir %s" % tempdir)
345
346     traffic.generate_users_and_groups(ldb,
347                                       opts.instance_id,
348                                       opts.fixed_password,
349                                       number_of_users,
350                                       opts.number_of_groups,
351                                       opts.group_memberships,
352                                       opts.max_members,
353                                       machine_accounts=len(conversations),
354                                       traffic_accounts=True)
355
356     accounts = traffic.generate_replay_accounts(ldb,
357                                                 opts.instance_id,
358                                                 len(conversations),
359                                                 opts.fixed_password)
360
361     statsdir = traffic.mk_masked_dir(tempdir, 'stats')
362
363     if opts.traffic_summary:
364         if opts.traffic_summary == '-':
365             summary_dest = sys.stdout
366         else:
367             summary_dest = open(opts.traffic_summary, 'w')
368
369         logger.info("Writing traffic summary")
370         summaries = []
371         for c in conversations:
372             summaries += c.replay_as_summary_lines()
373
374         summaries.sort()
375         for (time, line) in summaries:
376             print(line, file=summary_dest)
377
378         exit(0)
379
380     traffic.replay(conversations, host,
381                    lp=lp,
382                    creds=creds,
383                    accounts=accounts,
384                    dns_rate=opts.dns_rate,
385                    duration=duration,
386                    badpassword_frequency=opts.badpassword_frequency,
387                    prefer_kerberos=opts.prefer_kerberos,
388                    statsdir=statsdir,
389                    domain=domain,
390                    base_dn=ldb.domain_dn(),
391                    ou=traffic.ou_name(ldb, opts.instance_id),
392                    tempdir=tempdir,
393                    domain_sid=ldb.get_domain_sid())
394
395     if opts.timing_data == '-':
396         timing_dest = sys.stdout
397     elif opts.timing_data is None:
398         timing_dest = None
399     else:
400         timing_dest = open(opts.timing_data, 'w')
401
402     logger.info("Generating statistics")
403     traffic.generate_stats(statsdir, timing_dest)
404
405     if not opts.preserve_tempdir:
406         logger.info("Removing temporary directory")
407         shutil.rmtree(tempdir)
408
409
410 main()