]>
git.ipfire.org Git - thirdparty/pdns.git/blob - regression-tests.common/proxyprotocol.py
7 class ProxyProtocol(object):
8 MAGIC
= b
'\x0D\x0A\x0D\x0A\x00\x0D\x0A\x51\x55\x49\x54\x0A'
9 # Header is magic + versioncommand (1) + family (1) + content length (2)
10 HEADER_SIZE
= len(MAGIC
) + 1 + 1 + 2
16 def parseHeader(self
, data
):
17 if len(data
) < self
.HEADER_SIZE
:
20 if data
[:len(self
.MAGIC
)] != self
.MAGIC
:
23 value
= struct
.unpack('!B', bytes(bytearray([data
[12]])))[0]
24 self
.version
= value
>> 4
25 if self
.version
!= 0x02:
28 self
.command
= value
& ~
0x20
30 self
.offset
= self
.HEADER_SIZE
32 if self
.command
== 0x00:
34 elif self
.command
== 0x01:
35 value
= struct
.unpack('!B', bytes(bytearray([data
[13]])))[0]
36 self
.family
= value
>> 4
37 if self
.family
== 0x01:
39 elif self
.family
== 0x02:
44 self
.protocol
= value
& ~
0xF0
45 if self
.protocol
== 0x01:
47 elif self
.protocol
== 0x02:
54 self
.contentLen
= struct
.unpack("!H", data
[14:16])[0]
57 if self
.contentLen
< (self
.addrSize
* 2 + self
.PORT_SIZE
* 2):
62 def getAddr(self
, data
):
63 if len(data
) < (self
.consumed() + self
.addrSize
):
67 if self
.family
== 0x01:
68 value
= socket
.inet_ntop(socket
.AF_INET
, data
[self
.offset
:self
.offset
+ self
.addrSize
])
70 value
= socket
.inet_ntop(socket
.AF_INET6
, data
[self
.offset
:self
.offset
+ self
.addrSize
])
72 self
.offset
= self
.offset
+ self
.addrSize
75 def getPort(self
, data
):
76 if len(data
) < (self
.consumed() + self
.PORT_SIZE
):
79 value
= struct
.unpack('!H', data
[self
.offset
:self
.offset
+ self
.PORT_SIZE
])[0]
80 self
.offset
= self
.offset
+ self
.PORT_SIZE
83 def parseAddressesAndPorts(self
, data
):
87 if len(data
) < (self
.consumed() + self
.addrSize
* 2 + self
.PORT_SIZE
* 2):
90 self
.source
= self
.getAddr(data
)
91 self
.destination
= self
.getAddr(data
)
92 self
.sourcePort
= self
.getPort(data
)
93 self
.destinationPort
= self
.getPort(data
)
96 def parseAdditionalValues(self
, data
):
101 if len(data
) < (self
.HEADER_SIZE
+ self
.contentLen
):
104 remaining
= self
.HEADER_SIZE
+ self
.contentLen
- self
.consumed()
105 if len(data
) < remaining
:
108 while remaining
>= 3:
109 valueType
= struct
.unpack("!B", bytes(bytearray([data
[self
.offset
]])))[0]
110 self
.offset
= self
.offset
+ 1
111 valueLen
= struct
.unpack("!H", data
[self
.offset
:self
.offset
+2])[0]
112 self
.offset
= self
.offset
+ 2
114 remaining
= remaining
- 3
116 if valueLen
> remaining
:
118 self
.values
.append([valueType
, data
[self
.offset
:self
.offset
+valueLen
]])
119 self
.offset
= self
.offset
+ valueLen
120 remaining
= remaining
- valueLen
123 self
.values
.append([valueType
, ""])
128 def getPayload(cls
, local
, tcp
, v6
, source
, destination
, sourcePort
, destinationPort
, values
):
129 payload
= copy
.deepcopy(cls
.MAGIC
)
137 value
= struct
.pack('!B', (version
<< 4) + command
)
138 payload
= payload
+ value
148 # sorry but compatibility with python 2 is awful for this,
149 # not going to waste time on it
157 value
= struct
.pack('!B', (family
<< 4) + protocol
)
158 payload
= payload
+ value
162 contentSize
= contentSize
+ addrSize
* 2 + cls
.PORT_SIZE
*2
166 valuesSize
= valuesSize
+ 3 + len(value
[1])
168 contentSize
= contentSize
+ valuesSize
170 value
= struct
.pack('!H', contentSize
)
171 payload
= payload
+ value
179 value
= socket
.inet_pton(af
, source
)
180 payload
= payload
+ value
181 value
= socket
.inet_pton(af
, destination
)
182 payload
= payload
+ value
183 value
= struct
.pack('!H', sourcePort
)
184 payload
= payload
+ value
185 value
= struct
.pack('!H', destinationPort
)
186 payload
= payload
+ value
189 valueType
= struct
.pack('!B', value
[0])
190 valueLen
= struct
.pack('!H', len(value
[1]))
191 payload
= payload
+ valueType
+ valueLen
+ value
[1]