python/samba/gp_parse: PY3 open file non-binary mode for write_binary
[amitay/samba.git] / python / samba / gp_parse / gp_csv.py
1 # GPO Parser for audit extensions
2 #
3 # Copyright (C) Andrew Bartlett <abartlet@samba.org> 2018
4 # Written by Garming Sam <garming@catalyst.net.nz>
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
20 import codecs
21 import csv
22 import io
23
24 from io import BytesIO
25 from xml.etree.ElementTree import Element, SubElement
26 from samba.compat import PY3
27 from samba.gp_parse import GPParser
28
29 # [MS-GPAC] Group Policy Audit Configuration
30 class GPAuditCsvParser(GPParser):
31     encoding = 'utf-8'
32     header = None
33     lines = []
34
35     def parse(self, contents):
36         self.lines = []
37         reader = UnicodeReader(BytesIO(contents),
38                                encoding=self.encoding)
39
40         self.header = reader.next()
41         for row in reader:
42             line = {}
43             for i, x in enumerate(row):
44                 line[self.header[i]] = x
45
46             self.lines.append(line)
47             # print line
48
49     def write_xml(self, filename):
50         with open(filename, 'wb') as f:
51             root = Element('CsvFile')
52             child = SubElement(root, 'Row')
53             for e in self.header:
54                 value = SubElement(child, 'Value')
55                 value.text = e
56
57             for line in self.lines:
58                 child = SubElement(root, 'Row')
59                 for e, title in [(line[x], x) for x in self.header]:
60                     value = SubElement(child, 'Value')
61                     value.text = e
62
63                     # Metadata for generalization
64                     if title == 'Policy Target' and e != '':
65                         value.attrib['user_id'] = 'TRUE'
66                     if (title == 'Setting Value' and e != '' and
67                         (line['Subcategory'] == 'RegistryGlobalSacl' or
68                          line['Subcategory'] == 'FileGlobalSacl')):
69                         value.attrib['acl'] = 'TRUE'
70
71             self.write_pretty_xml(root, f)
72
73
74         # contents = codecs.open(filename, encoding='utf-8').read()
75         # self.load_xml(fromstring(contents))
76
77     def load_xml(self, root):
78         header = True
79         self.lines = []
80
81         for r in root.findall('Row'):
82             if header:
83                 header = False
84                 self.header = []
85                 for v in r.findall('Value'):
86                     self.header.append(v.text.decode(self.output_encoding))
87             else:
88                 line = {}
89                 for i, v in enumerate(r.findall('Value')):
90                     line[self.header[i]] = v.text if v.text is not None else ''
91                     line[self.header[i]] = line[self.header[i]].decode(self.output_encoding)
92
93                 self.lines.append(line)
94
95     def write_binary(self, filename):
96         from io import open
97         with open(filename, 'w', self.encoding) as f:
98             # In this case "binary" means "utf-8", so we let Python do that.
99             writer = csv.writer(f, quoting=csv.QUOTE_MINIMAL)
100             writer.writerow(self.header)
101             for line in self.lines:
102                 writer.writerow([line[x] for x in self.header])
103
104
105 # The following classes come from the Python documentation
106 # https://docs.python.org/3.0/library/csv.html
107
108
109 class UTF8Recoder:
110     """
111     Iterator that reads an encoded stream and reencodes the input to UTF-8
112     """
113     def __init__(self, f, encoding):
114         self.reader = codecs.getreader(encoding)(f)
115
116     def __iter__(self):
117         return self
118
119     def next(self):
120         return next(self.reader).encode("utf-8")
121
122     __next__ = next
123
124 class UnicodeReader:
125     """
126     A CSV reader which will iterate over lines in the CSV file "f",
127     which is encoded in the given encoding.
128     """
129
130     def __init__(self, f, dialect=csv.excel, encoding="utf-8", **kwds):
131         f = UTF8Recoder(f, encoding)
132         self.reader = csv.reader(f, dialect=dialect, **kwds)
133
134     def next(self):
135         row = next(self.reader)
136         return [unicode(s, "utf-8") for s in row]
137
138     def __iter__(self):
139         return self
140
141     __next__ = next
142
143 class UnicodeWriter:
144     """
145     A CSV writer which will write rows to CSV file "f",
146     which is encoded in the given encoding.
147     """
148
149     def __init__(self, f, dialect=csv.excel, encoding="utf-8", **kwds):
150         # Redirect output to a queue
151         self.queue = io.StringIO()
152         self.writer = csv.writer(self.queue, dialect=dialect, **kwds)
153         self.stream = f
154         self.encoder = codecs.getincrementalencoder(encoding)()
155
156     def writerow(self, row):
157         self.writer.writerow([s.encode("utf-8") for s in row])
158         # Fetch UTF-8 output from the queue ...
159         data = self.queue.getvalue()
160         data = data.decode("utf-8")
161         # ... and reencode it into the target encoding
162         data = self.encoder.encode(data)
163         # write to the target stream
164         self.stream.write(data)
165         # empty queue
166         self.queue.truncate(0)
167
168     def writerows(self, rows):
169         for row in rows:
170             self.writerow(row)