修改为东南天坐标系
This commit is contained in:
@@ -0,0 +1,399 @@
|
||||
import enum
|
||||
import typing
|
||||
|
||||
from ._types import StrOrBytes
|
||||
from .compat import singledispatchmethod
|
||||
from .exceptions import ProtocolError
|
||||
from .utils import (
|
||||
AddressType,
|
||||
decode_address,
|
||||
encode_address,
|
||||
get_address_port_tuple_from_address,
|
||||
)
|
||||
|
||||
|
||||
class SOCKS5AuthMethod(bytes, enum.Enum):
|
||||
"""Enumeration of SOCKS5 authentication methods."""
|
||||
|
||||
NO_AUTH_REQUIRED = b"\x00"
|
||||
GSSAPI = b"\x01"
|
||||
USERNAME_PASSWORD = b"\x02"
|
||||
NO_ACCEPTABLE_METHODS = b"\xFF"
|
||||
|
||||
|
||||
class SOCKS5Command(bytes, enum.Enum):
|
||||
"""Enumeration of SOCKS5 commands."""
|
||||
|
||||
CONNECT = b"\x01"
|
||||
BIND = b"\x02"
|
||||
UDP_ASSOCIATE = b"\x03"
|
||||
|
||||
|
||||
class SOCKS5AType(bytes, enum.Enum):
|
||||
"""Enumeration of SOCKS5 address types."""
|
||||
|
||||
IPV4_ADDRESS = b"\x01"
|
||||
DOMAIN_NAME = b"\x03"
|
||||
IPV6_ADDRESS = b"\x04"
|
||||
|
||||
@classmethod
|
||||
def from_atype(cls, atype: AddressType) -> "SOCKS5AType":
|
||||
if atype == AddressType.IPV4:
|
||||
return SOCKS5AType.IPV4_ADDRESS
|
||||
elif atype == AddressType.DN:
|
||||
return SOCKS5AType.DOMAIN_NAME
|
||||
elif atype == AddressType.IPV6:
|
||||
return SOCKS5AType.IPV6_ADDRESS
|
||||
raise ValueError(atype)
|
||||
|
||||
|
||||
class SOCKS5ReplyCode(bytes, enum.Enum):
|
||||
"""Enumeration of SOCKS5 reply codes."""
|
||||
|
||||
SUCCEEDED = b"\x00"
|
||||
GENERAL_SERVER_FAILURE = b"\x01"
|
||||
CONNECTION_NOT_ALLOWED_BY_RULESET = b"\x02"
|
||||
NETWORK_UNREACHABLE = b"\x03"
|
||||
HOST_UNREACHABLE = b"\x04"
|
||||
CONNECTION_REFUSED = b"\x05"
|
||||
TTL_EXPIRED = b"\x06"
|
||||
COMMAND_NOT_SUPPORTED = b"\x07"
|
||||
ADDRESS_TYPE_NOT_SUPPORTED = b"\x08"
|
||||
|
||||
|
||||
class SOCKS5AuthMethodsRequest(typing.NamedTuple):
|
||||
"""Encapsulates a request to the proxy for available authentication methods.
|
||||
|
||||
Args:
|
||||
methods: A list of acceptable authentication methods.
|
||||
"""
|
||||
|
||||
methods: typing.List[SOCKS5AuthMethod]
|
||||
|
||||
def dumps(self) -> bytes:
|
||||
"""Packs the instance into a raw binary in the appropriate form."""
|
||||
|
||||
return b"".join(
|
||||
[
|
||||
b"\x05",
|
||||
len(self.methods).to_bytes(1, byteorder="big"),
|
||||
b"".join(self.methods),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
class SOCKS5AuthReply(typing.NamedTuple):
|
||||
"""Encapsulates a reply from the proxy with the authentication method to be used.
|
||||
|
||||
Args:
|
||||
method: The authentication method to be used.
|
||||
|
||||
Raises:
|
||||
ProtocolError: If the data does not conform with the expected structure.
|
||||
"""
|
||||
|
||||
method: SOCKS5AuthMethod
|
||||
|
||||
@classmethod
|
||||
def loads(cls, data: bytes) -> "SOCKS5AuthReply":
|
||||
"""Unpacks the authentication reply data into an instance.
|
||||
|
||||
Returns:
|
||||
The unpacked authentication reply instance.
|
||||
|
||||
Raises:
|
||||
ProtocolError: If the data does not match the spec.
|
||||
"""
|
||||
if len(data) != 2:
|
||||
raise ProtocolError("Malformed reply")
|
||||
|
||||
try:
|
||||
return cls(method=SOCKS5AuthMethod(data[1:2]))
|
||||
except ValueError as exc:
|
||||
raise ProtocolError("Malformed reply") from exc
|
||||
|
||||
|
||||
class SOCKS5UsernamePasswordRequest(typing.NamedTuple):
|
||||
"""Encapsulates a username/password authentication request to the proxy server."""
|
||||
|
||||
username: bytes
|
||||
password: bytes
|
||||
|
||||
def dumps(self) -> bytes:
|
||||
"""Packs the instance into a raw binary in the appropriate form.
|
||||
|
||||
Returns:
|
||||
The packed request.
|
||||
"""
|
||||
return b"".join(
|
||||
[
|
||||
b"\x01",
|
||||
len(self.username).to_bytes(1, byteorder="big"),
|
||||
self.username,
|
||||
len(self.password).to_bytes(1, byteorder="big"),
|
||||
self.password,
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
class SOCKS5UsernamePasswordReply(typing.NamedTuple):
|
||||
"""Encapsulates a username/password authentication reply from the proxy server."""
|
||||
|
||||
success: bool
|
||||
|
||||
@classmethod
|
||||
def loads(cls, data: bytes) -> "SOCKS5UsernamePasswordReply":
|
||||
"""Unpacks the reply authentication data into an instance.
|
||||
|
||||
Returns:
|
||||
The unpacked authentication reply instance.
|
||||
"""
|
||||
return cls(success=data == b"\x01\x00")
|
||||
|
||||
|
||||
class SOCKS5CommandRequest(typing.NamedTuple):
|
||||
"""Encapsulates a command request to the proxy server.
|
||||
|
||||
Args:
|
||||
command: The command to request.
|
||||
atype: The address type of the addr field.
|
||||
addr: Address of the target host.
|
||||
port: The port number to connect to on the target host.
|
||||
"""
|
||||
|
||||
command: SOCKS5Command
|
||||
atype: SOCKS5AType
|
||||
addr: bytes
|
||||
port: int
|
||||
|
||||
@classmethod
|
||||
def from_address(
|
||||
cls,
|
||||
command: SOCKS5Command,
|
||||
address: typing.Union[StrOrBytes, typing.Tuple[StrOrBytes, int]],
|
||||
) -> "SOCKS5CommandRequest":
|
||||
"""Convenience class method to build an instance from command and address.
|
||||
|
||||
Args:
|
||||
command: The command to request.
|
||||
address: A string in the form 'HOST:PORT' or a tuple of ip address string
|
||||
and port number. The address type will be inferred.
|
||||
|
||||
Returns:
|
||||
A SOCKS5CommandRequest instance.
|
||||
|
||||
Raises:
|
||||
SOCKSError: If a domain name or IPv6 address was supplied.
|
||||
"""
|
||||
address, port = get_address_port_tuple_from_address(address)
|
||||
atype, encoded_addr = encode_address(address)
|
||||
return cls(
|
||||
command=command,
|
||||
atype=SOCKS5AType.from_atype(atype),
|
||||
addr=encoded_addr,
|
||||
port=port,
|
||||
)
|
||||
|
||||
def dumps(self) -> bytes:
|
||||
"""Packs the instance into a raw binary in the appropriate form.
|
||||
|
||||
Returns:
|
||||
The packed request.
|
||||
"""
|
||||
return b"".join(
|
||||
[
|
||||
b"\x05",
|
||||
self.command,
|
||||
b"\x00",
|
||||
self.atype,
|
||||
self.packed_addr,
|
||||
(self.port).to_bytes(2, byteorder="big"),
|
||||
]
|
||||
)
|
||||
|
||||
@property
|
||||
def packed_addr(self) -> bytes:
|
||||
"""Property returning the packed address in the correct form for its type."""
|
||||
if self.atype == SOCKS5AType.IPV4_ADDRESS:
|
||||
assert len(self.addr) == 4
|
||||
return self.addr
|
||||
elif self.atype == SOCKS5AType.IPV6_ADDRESS:
|
||||
assert len(self.addr) == 16
|
||||
return self.addr
|
||||
else:
|
||||
length = len(self.addr)
|
||||
return length.to_bytes(1, byteorder="big") + self.addr
|
||||
|
||||
|
||||
class SOCKS5Reply(typing.NamedTuple):
|
||||
"""Encapsulates a reply from the SOCKS5 proxy server
|
||||
|
||||
Args:
|
||||
reply_code: The code representing the type of reply.
|
||||
atype: The address type of the addr field.
|
||||
addr: Optional IP address returned.
|
||||
port: The port number returned.
|
||||
"""
|
||||
|
||||
reply_code: SOCKS5ReplyCode
|
||||
atype: SOCKS5AType
|
||||
addr: str
|
||||
port: int
|
||||
|
||||
@classmethod
|
||||
def loads(cls, data: bytes) -> "SOCKS5Reply":
|
||||
"""Unpacks the reply data into an instance.
|
||||
|
||||
Returns:
|
||||
The unpacked reply instance.
|
||||
|
||||
Raises:
|
||||
ProtocolError: If the data does not match the spec.
|
||||
"""
|
||||
if data[0:1] != b"\x05":
|
||||
raise ProtocolError("Malformed reply")
|
||||
|
||||
try:
|
||||
atype = SOCKS5AType(data[3:4])
|
||||
|
||||
return cls(
|
||||
reply_code=SOCKS5ReplyCode(data[1:2]),
|
||||
atype=atype,
|
||||
addr=decode_address(AddressType.from_socks5_atype(atype), data[4:-2]),
|
||||
port=int.from_bytes(data[-2:], byteorder="big"),
|
||||
)
|
||||
except ValueError as exc:
|
||||
raise ProtocolError("Malformed reply") from exc
|
||||
|
||||
|
||||
class SOCKS5Datagram(typing.NamedTuple):
|
||||
"""Encapsulates a SOCKS5 datagram for UDP connections.
|
||||
|
||||
Currently not implemented.
|
||||
"""
|
||||
|
||||
atype: SOCKS5AType
|
||||
addr: bytes
|
||||
port: int
|
||||
data: bytes
|
||||
|
||||
fragment: int
|
||||
last_fragment: bool
|
||||
|
||||
@classmethod
|
||||
def loads(cls, data: bytes) -> "SOCKS5Datagram":
|
||||
raise NotImplementedError() # pragma: nocover
|
||||
|
||||
def dumps(self) -> bytes:
|
||||
raise NotImplementedError() # pragma: nocover
|
||||
|
||||
|
||||
class SOCKS5State(enum.IntEnum):
|
||||
"""Enumeration of SOCKS5 protocol states."""
|
||||
|
||||
CLIENT_AUTH_REQUIRED = 1
|
||||
SERVER_AUTH_REPLY = 2
|
||||
CLIENT_AUTHENTICATED = 3
|
||||
TUNNEL_READY = 4
|
||||
CLIENT_WAITING_FOR_USERNAME_PASSWORD = 5
|
||||
SERVER_VERIFY_USERNAME_PASSWORD = 6
|
||||
MUST_CLOSE = 7
|
||||
|
||||
|
||||
SOCKS5RequestType = typing.Union[SOCKS5AuthMethodsRequest, SOCKS5CommandRequest]
|
||||
|
||||
|
||||
class SOCKS5Connection:
|
||||
"""Encapsulates a SOCKS5 connection.
|
||||
|
||||
Packs request objects into data suitable to be send and unpacks reply
|
||||
data into their appropriate reply objects.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._data_to_send = bytearray()
|
||||
self._received_data = bytearray()
|
||||
self._state = SOCKS5State.CLIENT_AUTH_REQUIRED
|
||||
|
||||
@property
|
||||
def state(self) -> SOCKS5State:
|
||||
"""Returns the current state of the protocol."""
|
||||
return self._state
|
||||
|
||||
@singledispatchmethod # type: ignore
|
||||
def send(self, request: SOCKS5RequestType) -> None:
|
||||
"""Packs a request object and adds it to the send data buffer.
|
||||
|
||||
Also progresses the protocol state of the connection.
|
||||
|
||||
Args:
|
||||
request: The request instance to be packed.
|
||||
"""
|
||||
raise NotImplementedError() # pragma: nocover
|
||||
|
||||
@send.register(SOCKS5AuthMethodsRequest) # type: ignore
|
||||
def _auth_methods(self, request: SOCKS5AuthMethodsRequest) -> None:
|
||||
self._data_to_send += request.dumps()
|
||||
self._state = SOCKS5State.SERVER_AUTH_REPLY
|
||||
|
||||
@send.register(SOCKS5UsernamePasswordRequest) # type: ignore
|
||||
def _auth_username_password(self, request: SOCKS5UsernamePasswordRequest) -> None:
|
||||
if self._state != SOCKS5State.CLIENT_WAITING_FOR_USERNAME_PASSWORD:
|
||||
raise ProtocolError("Not currently waiting for username and password")
|
||||
self._state = SOCKS5State.SERVER_VERIFY_USERNAME_PASSWORD
|
||||
self._data_to_send += request.dumps()
|
||||
|
||||
@send.register(SOCKS5CommandRequest) # type: ignore
|
||||
def _command(self, request: SOCKS5AuthMethodsRequest) -> None:
|
||||
if self._state < SOCKS5State.CLIENT_AUTHENTICATED:
|
||||
raise ProtocolError(
|
||||
"SOCKS5 connections must be authenticated before sending a request"
|
||||
)
|
||||
self._data_to_send += request.dumps()
|
||||
|
||||
def receive_data(
|
||||
self, data: bytes
|
||||
) -> typing.Union[SOCKS5AuthReply, SOCKS5Reply, SOCKS5UsernamePasswordReply]:
|
||||
"""Unpacks response data into a reply object.
|
||||
|
||||
Args:
|
||||
data: The raw response data from the proxy server.
|
||||
|
||||
Returns:
|
||||
A reply instance corresponding to the connection state and reply data.
|
||||
"""
|
||||
if self._state == SOCKS5State.SERVER_AUTH_REPLY:
|
||||
auth_reply = SOCKS5AuthReply.loads(data)
|
||||
if auth_reply.method == SOCKS5AuthMethod.USERNAME_PASSWORD:
|
||||
self._state = SOCKS5State.CLIENT_WAITING_FOR_USERNAME_PASSWORD
|
||||
elif auth_reply.method == SOCKS5AuthMethod.NO_AUTH_REQUIRED:
|
||||
self._state = SOCKS5State.CLIENT_AUTHENTICATED
|
||||
return auth_reply
|
||||
|
||||
if self._state == SOCKS5State.SERVER_VERIFY_USERNAME_PASSWORD:
|
||||
username_password_reply = SOCKS5UsernamePasswordReply.loads(data)
|
||||
if username_password_reply.success:
|
||||
self._state = SOCKS5State.CLIENT_AUTHENTICATED
|
||||
else:
|
||||
self._state = SOCKS5State.MUST_CLOSE
|
||||
return username_password_reply
|
||||
|
||||
if self._state == SOCKS5State.CLIENT_AUTHENTICATED:
|
||||
reply = SOCKS5Reply.loads(data)
|
||||
if reply.reply_code == SOCKS5ReplyCode.SUCCEEDED:
|
||||
self._state = SOCKS5State.TUNNEL_READY
|
||||
else:
|
||||
self._state = SOCKS5State.MUST_CLOSE
|
||||
|
||||
return reply
|
||||
|
||||
raise NotImplementedError() # pragma: nocover
|
||||
|
||||
def data_to_send(self) -> bytes:
|
||||
"""Returns the data to be sent via the I/O library of choice.
|
||||
|
||||
Also clears the connection's buffer.
|
||||
"""
|
||||
data = bytes(self._data_to_send)
|
||||
self._data_to_send = bytearray()
|
||||
return data
|
||||
Reference in New Issue
Block a user