1# Copyright (c) 2003-2016 CORE Security Technologies
2#
3# This software is provided under under a slightly modified version
4# of the Apache Software License. See the accompanying LICENSE file
5# for more information.
6#
7# Description:
8#  WPS packets
9#
10# Author:
11# Aureliano Calvo
12
13
14import array
15import struct
16
17from impacket.helper import ProtocolPacket, Byte, Bit
18
19
20class ArrayBuilder(object):
21
22    def from_ary(self, ary):
23        return ary
24
25    def to_ary(self, value):
26        return array.array("B", value)
27
28class ByteBuilder(object):
29
30    def from_ary(self, ary):
31        return ary[0]
32
33    def to_ary(self, value):
34        return array.array('B', [value])
35
36class StringBuilder(object):
37    def from_ary(self, ary):
38        return ary.tostring()
39
40    def to_ary(self, value):
41        return array.array('B', value)
42
43class NumBuilder(object):
44    """Converts back and forth between arrays and numbers in network byte-order"""
45
46    def __init__(self, size):
47        """size: number of bytes in the field"""
48        self.size = size
49
50    def from_ary(self, ary):
51        if len(ary) != self.size:
52            raise Exception("Expected %s size but got %s" % (self.size, len(ary)))
53        return reduce( lambda ac, x: ac * 256 + x, ary, 0)
54
55    def to_ary(self, value0):
56        value = value0
57        rv = array.array('B')
58        for _ in xrange(self.size):
59            value, mod = divmod(value, 256)
60            rv.append(mod)
61
62        if value != 0:
63            raise Exception("%s is too big. Max size: %s" % (value0, self.size))
64
65        rv.reverse()
66        return rv
67
68class TLVContainer(object):
69
70    def builder(self, kind):
71        return self.builders.get(kind, self.default_builder)
72
73    def from_ary(self, ary):
74        i = 0
75        while i<len(ary):
76            kind = self.ary2n(ary, i)
77            length = self.ary2n(ary, i+2)
78            i+=4
79            value = ary[i:i+length]
80            self.elems.append((kind, value))
81            i += length
82
83        return self
84
85    def __init__(self, builders, default_builder = ArrayBuilder(), descs=None):
86        self.builders = builders
87        self.default_builder = default_builder
88        self.elems = []
89        self.descs = descs or {}
90
91    def append(self, kind, value):
92        self.elems.append((kind, self.builder(kind).to_ary(value)))
93
94    def __iter__(self):
95        return ((k, self.builder(k).from_ary(v)) for k,v in self.elems)
96
97    def all(self, kind):
98        return [e[1] for e in self if e[0] == kind]
99
100    def __contains__(self, kind):
101        return len(self.all(kind)) != 0
102
103    def first(self, kind):
104        return self.all(kind)[0]
105
106    def to_ary(self):
107        ary = array.array('B')
108        for k,v in self.elems:
109            ary.extend(self.n2ary(k))
110            ary.extend(self.n2ary(len(v)))
111            ary.extend(v)
112
113        return ary
114
115
116    def get_packet(self):
117        return self.to_ary().tostring()
118
119    def set_parent(self, my_parent):
120        self.__parent = my_parent
121
122    def parent(self):
123        return self.__parent
124
125    def n2ary(self, n):
126        return array.array("B", struct.pack(">H",n))
127
128    def ary2n(self, ary, i=0):
129        return struct.unpack(">H", ary[i:i+2].tostring())[0]
130
131    def __repr__(self):
132        def desc(kind):
133            return self.descs[kind] if kind in self.descs else kind
134
135        return "<TLVContainer %s>" % repr([(desc(k), self.builder(k).from_ary(v)) for (k,v) in self.elems])
136
137    def child(self):
138        return None
139
140class SCElem(object):
141    #Data elements as defined in section 11 of the WPS 1.0h spec.
142
143    AP_CHANNEL = 0x1001
144    ASSOCIATION_STATE = 0x1002
145    AUTHENTICATION_TYPE = 0x1003
146    AUTHENTICATION_TYPE_FLAGS = 0x1004
147    AUTHENTICATOR = 0x1005
148    CONFIG_METHODS = 0x1008
149    CONFIGURATION_ERROR = 0x1009
150    CONFIRMATION_URL4 = 0x100A
151    CONFIRMATION_URL6 = 0x100B
152    CONNECTION_TYPE = 0X100C
153    CONNECTION_TYPE_FLAGS = 0X100D
154    CREDENTIAL = 0X100E
155    DEVICE_NAME = 0x1011
156    DEVICE_PASSWORD_ID = 0x1012
157    E_HASH1 = 0x1014
158    E_HASH2 = 0x1015
159    E_SNONCE1 = 0x1016
160    E_SNONCE2 = 0x1017
161    ENCRYPTED_SETTINGS = 0x1018
162    ENCRYPTION_TYPE = 0X100F
163    ENCRYPTION_TYPE_FLAGS = 0x1010
164    ENROLLEE_NONCE = 0x101A
165    FEATURE_ID = 0x101B
166    IDENTITY = 0X101C
167    INDENTITY_PROOF = 0X101D
168    KEY_WRAP_AUTHENTICATOR = 0x101E
169    KEY_IDENTIFIER = 0X101F
170    MAC_ADDRESS = 0x1020
171    MANUFACTURER = 0x1021
172    MESSAGE_TYPE = 0x1022
173    MODEL_NAME = 0x1023
174    MODEL_NUMBER = 0x1024
175    NETWORK_INDEX = 0x1026
176    NETWORK_KEY = 0x1027
177    NETWORK_KEY_INDEX = 0x1028
178    NEW_DEVICE_NAME = 0x1029
179    NEW_PASSWORD = 0x102A
180    OOB_DEVICE_PASSWORD = 0X102C
181    OS_VERSION= 0X102D
182    POWER_LEVEL = 0X102F
183    PSK_CURRENT = 0x1030
184    PSK_MAX = 0x1031
185    PUBLIC_KEY = 0x1032
186    RADIO_ENABLED = 0x1033
187    REBOOT = 0x1034
188    REGISTRAR_CURRENT = 0x1035
189    REGISTRAR_ESTABLISHED = 0x1036
190    REGISTRAR_LIST = 0x1037
191    REGISTRAR_MAX = 0x1038
192    REGISTRAR_NONCE = 0x1039
193    REQUEST_TYPE = 0x103A
194    RESPONSE_TYPE = 0x103B
195    RF_BANDS = 0X103C
196    R_HASH1 = 0X103D
197    R_HASH2 = 0X103E
198    R_SNONCE1 = 0X103F
199    R_SNONCE2 = 0x1040
200    SELECTED_REGISTRAR = 0x1041
201    SERIAL_NUMBER = 0x1042
202    WPS_STATE = 0x1044
203    SSID = 0x1045
204    TOTAL_NETWORKS = 0x1046
205    UUID_E = 0x1047
206    UUID_R = 0x1048
207    VENDOR_EXTENSION = 0x1049
208    VERSION = 0x104A
209    X_509_CERTIFICATE_REQUEST = 0x104B
210    X_509_CERTIFICATE = 0x104C
211    EAP_IDENTITY = 0x104D
212    MESSAGE_COUNTER = 0x104E
213    PUBLIC_KEY_HASH = 0x104F
214    REKEY_KEY = 0x1050
215    KEY_LIFETIME = 0x1051
216    PERMITTED_CONFIG_METHODS = 0x1052
217    SELECTED_REGISTRAR_CONFIG_METHODS= 0x1053
218    PRIMARY_DEVICE_TYPE = 0x1054
219    SECONDARY_DEVICE_TYPE_LIST = 0x1055
220    PORTABLE_DEVICE = 0x1056
221    AP_SETUP_LOCKED = 0x1057
222    APPLICATION_EXTENSION = 0x1058
223    EAP_TYPE = 0x1059
224    INITIALIZATION_VECTOR = 0x1060
225    KEY_PROVIDED_AUTOMATICALLY = 0x1061
226    _802_1X_ENABLED = 0x1062
227    APP_SESSION_KEY = 0x1063
228    WEP_TRANSMIT_KEY = 0x1064
229
230class MessageType(object):
231    """Message types according to WPS 1.0h spec, section 11"""
232
233    BEACON = 0x01
234    PROBE_REQUEST = 0x02
235    PROBE_RESPONSE = 0x03
236    M1 = 0x04
237    M2 = 0x05
238    M2D = 0x06
239    M3 = 0x07
240    M4 = 0x08
241    M5 = 0x09
242    M6 = 0x0A
243    M7 = 0x0B
244    M8 = 0x0C
245    WSC_ACK = 0x0D
246    WSC_NACK = 0x0E
247    WSC_DONE = 0x0F
248
249class AuthTypeFlag(object):
250    OPEN = 0x0001
251    WPAPSK = 0x0002
252    SHARED = 0x0004
253    WPA = 0x0008
254    WPA2 = 0x0010
255    WPA2PSK = 0x0020
256
257AuthTypeFlag_ALL = AuthTypeFlag.OPEN | \
258        AuthTypeFlag.WPAPSK | \
259        AuthTypeFlag.SHARED | \
260        AuthTypeFlag.WPA | \
261        AuthTypeFlag.WPA2 | \
262        AuthTypeFlag.WPA2PSK
263
264class EncryptionTypeFlag(object):
265    NONE = 0x0001
266    WEP = 0x0002
267    TKIP = 0x0004
268    AES = 0x0008
269
270EncryptionTypeFlag_ALL = EncryptionTypeFlag.NONE | EncryptionTypeFlag.WEP | EncryptionTypeFlag.TKIP | EncryptionTypeFlag.AES
271
272class ConnectionTypeFlag(object):
273    ESS = 0x01
274    IBSS = 0x02
275
276class ConfigMethod(object):
277    USBA = 0x0001
278    ETHERNET = 0x0002
279    LABEL = 0x0004
280    DISPLAY = 0x0008
281    EXT_NFC_TOKEN = 0x0010
282    INT_NFC_TOKEN = 0x0020
283    NFC_INTERFACE = 0x0040
284    PUSHBUTTON = 0x0080
285    KEYPAD = 0x0100
286
287
288class OpCode(object):
289    WSC_START = 0x01
290    WSC_ACK = 0x02
291    WSC_NACK = 0x03
292    WSC_MSG = 0x04
293    WSC_DONE = 0x05
294    WSC_FRAG_ACK = 0x06
295
296class AssocState(object):
297    NOT_ASSOC = 0
298    CONN_SUCCESS = 1
299    CFG_FAILURE = 2
300    FAILURE = 3,
301    IP_FAILURE = 4
302
303class ConfigError(object):
304    NO_ERROR = 0
305    OOB_IFACE_READ_ERROR = 1
306    DECRYPTION_CRC_FAILURE = 2
307    _24_CHAN_NOT_SUPPORTED = 3
308    _50_CHAN_NOT_SUPPORTED = 4
309    SIGNAL_TOO_WEAK = 5
310    NETWORK_AUTH_FAILURE = 6
311    NETWORK_ASSOC_FAILURE = 7
312    NO_DHCP_RESPONSE = 8
313    FAILED_DHCP_CONFIG = 9
314    IP_ADDR_CONFLICT = 10
315    NO_CONN_TO_REGISTRAR = 11
316    MULTIPLE_PBC_DETECTED = 12
317    ROGUE_SUSPECTED = 13
318    DEVICE_BUSY = 14
319    SETUP_LOCKED = 15
320    MSG_TIMEOUT = 16
321    REG_SESS_TIMEOUT = 17
322    DEV_PASSWORD_AUTH_FAILURE = 18
323
324class DevicePasswordId(object):
325    DEFAULT = 0x0000
326    USER_SPECIFIED = 0x0001
327    MACHINE_SPECIFIED = 0x0002
328    REKEY = 0x0003
329    PUSHBUTTON = 0x0004
330    REGISTRAR_SPECIFIED = 0x0005
331
332class WpsState(object):
333    NOT_CONFIGURED = 0x01
334    CONFIGURED = 0x02
335
336
337class SimpleConfig(ProtocolPacket):
338    "For now, it supports Simple configs with the bits more_fragments and length_field not set"
339
340    header_size = 2
341    tail_size = 0
342
343    op_code = Byte(0)
344    flags = Byte(1)
345    more_fragments = Bit(1, 0)
346    length_field = Bit(1,1)
347
348    BUILDERS = {
349        SCElem.CONNECTION_TYPE: ByteBuilder(),
350        SCElem.CONNECTION_TYPE_FLAGS: ByteBuilder(),
351        SCElem.VERSION: ByteBuilder(),
352        SCElem.MESSAGE_TYPE: ByteBuilder(),
353        SCElem.NETWORK_INDEX: ByteBuilder(),
354        SCElem.NETWORK_KEY_INDEX: ByteBuilder(),
355        SCElem.POWER_LEVEL: ByteBuilder(),
356        SCElem.PSK_CURRENT: ByteBuilder(),
357        SCElem.PSK_MAX: ByteBuilder(),
358        SCElem.REGISTRAR_CURRENT: ByteBuilder(),
359        SCElem.REGISTRAR_MAX: ByteBuilder(),
360        SCElem.REQUEST_TYPE: ByteBuilder(),
361        SCElem.RESPONSE_TYPE: ByteBuilder(),
362        SCElem.RF_BANDS: ByteBuilder(),
363        SCElem.WPS_STATE: ByteBuilder(),
364        SCElem.TOTAL_NETWORKS: ByteBuilder(),
365        SCElem.VERSION: ByteBuilder(),
366        SCElem.WEP_TRANSMIT_KEY: ByteBuilder(),
367
368        SCElem.CONFIRMATION_URL4: StringBuilder(),
369        SCElem.CONFIRMATION_URL6: StringBuilder(),
370        SCElem.DEVICE_NAME: StringBuilder(),
371        SCElem.IDENTITY: StringBuilder(),
372        SCElem.MANUFACTURER: StringBuilder(),
373        SCElem.MODEL_NAME: StringBuilder(),
374        SCElem.MODEL_NUMBER: StringBuilder(),
375        SCElem.NEW_DEVICE_NAME: StringBuilder(),
376        SCElem.NEW_PASSWORD: StringBuilder(),
377        SCElem.SERIAL_NUMBER: StringBuilder(),
378        SCElem.EAP_IDENTITY: StringBuilder(),
379        SCElem.NETWORK_KEY: StringBuilder(),
380
381        SCElem.AP_CHANNEL: NumBuilder(2),
382        SCElem.ASSOCIATION_STATE: NumBuilder(2),
383        SCElem.AUTHENTICATION_TYPE: NumBuilder(2),
384        SCElem.AUTHENTICATION_TYPE_FLAGS: NumBuilder(2),
385        SCElem.CONFIG_METHODS: NumBuilder(2),
386        SCElem.CONFIGURATION_ERROR: NumBuilder(2),
387        SCElem.DEVICE_PASSWORD_ID: NumBuilder(2),
388        SCElem.ENCRYPTION_TYPE: NumBuilder(2),
389        SCElem.ENCRYPTION_TYPE_FLAGS: NumBuilder(2),
390        SCElem.MESSAGE_COUNTER: NumBuilder(8),
391        SCElem.KEY_LIFETIME: NumBuilder(4),
392        SCElem.PERMITTED_CONFIG_METHODS: NumBuilder(2),
393        SCElem.SELECTED_REGISTRAR_CONFIG_METHODS: NumBuilder(2),
394        SCElem.PUBLIC_KEY: NumBuilder(192),
395
396    }
397
398    @classmethod
399    def build_tlv_container(cls):
400        return TLVContainer(
401            builders=SimpleConfig.BUILDERS,
402            descs = dict( (v,k) for (k,v) in SCElem.__dict__.iteritems() )
403        )
404
405