Current File : //lib/python3/dist-packages/socksio/socks5.py
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
Mini Shell