-
Notifications
You must be signed in to change notification settings - Fork 5
/
combine-tailoring.py
executable file
·128 lines (105 loc) · 4.17 KB
/
combine-tailoring.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
#!/usr/bin/python3
# Copyright 2016 Red Hat Inc., Durham, North Carolina.
# All Rights Reserved.
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
#
# Authors:
# Martin Preisler <[email protected]>
# Birol Bilgin <[email protected]>
from __future__ import print_function
import argparse
import sys
from io import BytesIO
try:
from xml.etree import cElementTree as ElementTree
except ImportError:
import cElementTree as ElementTree
XCCDF12_NS = "http://checklists.nist.gov/xccdf/1.2"
def main():
parser = argparse.ArgumentParser(
description="Combine XCCDF or datastream and a tailoring file to form "
"just one file that contains all the profiles."
)
parser.add_argument(
"SCAP_INPUT", type=argparse.FileType("r"),
help="XCCDF or Source DataStream"
)
parser.add_argument(
"TAILORING_FILE", type=argparse.FileType("r"),
help="XCCDF 1.2 Tailoring file to insert"
)
parser.add_argument(
"--output", type=argparse.FileType("wb"), required=False,
default=None,
help="Resulting XCCDF or Source DataStream"
)
args = parser.parse_args()
input_tree = ElementTree.ElementTree()
input_tree.parse(args.SCAP_INPUT)
input_root = input_tree.getroot()
tailoring_tree = ElementTree.ElementTree()
tailoring_tree.parse(args.TAILORING_FILE)
tailoring_root = tailoring_tree.getroot()
benchmarks = list(input_root.findall(".//{%s}Benchmark" % (XCCDF12_NS)))
if len(benchmarks) == 0:
print(
"There is no Benchmark elements in input file '%s'" % (args.SCAP_INPUT.name),
file=sys.stderr)
sys.exit(1)
t_profiles = tailoring_root.findall(".//{%s}Profile" % (XCCDF12_NS))
if len(t_profiles) == 0:
print(
"There is no Profile elements in the tailored file '%s'" % (args.TAILORING_FILE.name),
file=sys.stderr)
sys.exit(1)
# As far as my tests goes you cannot have a tailored file that has
# profiles belongs to different benchmark checklists
# we just need to figure out which benchmark tailored profiles belongs to
b_index = -1
extended_profile = t_profiles[0].get("extends")
benchmark = None
profile_insert_point = None
for i, bench in enumerate(benchmarks):
if b_index != -1:
break
for profile in bench.findall("./{%s}Profile" % (XCCDF12_NS)):
if b_index == -1 and \
(profile.get("id") == extended_profile or
extended_profile is None):
b_index = i
benchmark = benchmarks[i]
profile_insert_point = profile
if profile_insert_point is None:
print(
"Couldn't find a suitable Profile insert point in the input file. "
"Please check the @extends attribute in your tailoring file to "
"make sure there is a matching profile in the input file. "
"Make sure there is at least one profile in the input file. It can "
"be empty or have just one select.",
file=sys.stderr)
sys.exit(1)
for profile_to_add in t_profiles:
index = list(benchmark).index(profile_insert_point)
benchmark.insert(index + 1, profile_to_add)
profile_insert_point = profile_to_add
if args.output:
input_tree.write(args.output)
else:
output = BytesIO()
input_tree.write(output)
sys.stdout.write(output.getvalue().decode("utf8"))
if __name__ == "__main__":
main()