Current File : //lib/python3/dist-packages/socksio/socks4.py
import enum
import typing
from ._types import StrOrBytes
from .exceptions import ProtocolError, SOCKSError
from .utils import (
AddressType,
decode_address,
encode_address,
get_address_port_tuple_from_address,
)
class SOCKS4ReplyCode(bytes, enum.Enum):
"""Enumeration of SOCKS4 reply codes."""
REQUEST_GRANTED = b"\x5A"
REQUEST_REJECTED_OR_FAILED = b"\x5B"
CONNECTION_FAILED = b"\x5C"
AUTHENTICATION_FAILED = b"\x5D"
class SOCKS4Command(bytes, enum.Enum):
"""Enumeration of SOCKS4 command codes."""
CONNECT = b"\x01"
BIND = b"\x02"
class SOCKS4Request(typing.NamedTuple):
"""Encapsulates a request to the SOCKS4 proxy server
Args:
command: The command to request.
port: The port number to connect to on the target host.
addr: IP address of the target host.
user_id: Optional user ID to be included in the request, if not supplied
the user *must* provide one in the packing operation.
"""
command: SOCKS4Command
port: int
addr: bytes
user_id: typing.Optional[bytes] = None
@classmethod
def from_address(
cls,
command: SOCKS4Command,
address: typing.Union[StrOrBytes, typing.Tuple[StrOrBytes, int]],
user_id: typing.Optional[bytes] = None,
) -> "SOCKS4Request":
"""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.
user_id: Optional user ID.
Returns:
A SOCKS4Request 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)
if atype != AddressType.IPV4:
raise SOCKSError(
"IPv6 addresses and domain names are not supported by SOCKS4"
)
return cls(command=command, addr=encoded_addr, port=port, user_id=user_id)
def dumps(self, user_id: typing.Optional[bytes] = None) -> bytes:
"""Packs the instance into a raw binary in the appropriate form.
Args:
user_id: Optional user ID as an override, if not provided the instance's
will be used, if none was provided at initialization an error is raised.
Returns:
The packed request.
Raises:
SOCKSError: If no user was specified in this call or on initialization.
"""
user_id = user_id or self.user_id
if user_id is None:
raise SOCKSError("SOCKS4 requires a user_id, none was specified")
return b"".join(
[
b"\x04",
self.command,
(self.port).to_bytes(2, byteorder="big"),
self.addr,
user_id,
b"\x00",
]
)
class SOCKS4ARequest(typing.NamedTuple):
"""Encapsulates a request to the SOCKS4A proxy server
Args:
command: The command to request.
port: The port number to connect to on the target host.
addr: IP address of the target host.
user_id: Optional user ID to be included in the request, if not supplied
the user *must* provide one in the packing operation.
"""
command: SOCKS4Command
port: int
addr: bytes
user_id: typing.Optional[bytes] = None
@classmethod
def from_address(
cls,
command: SOCKS4Command,
address: typing.Union[StrOrBytes, typing.Tuple[StrOrBytes, int]],
user_id: typing.Optional[bytes] = None,
) -> "SOCKS4ARequest":
"""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.
user_id: Optional user ID.
Returns:
A SOCKS4ARequest instance.
"""
address, port = get_address_port_tuple_from_address(address)
atype, encoded_addr = encode_address(address)
return cls(command=command, addr=encoded_addr, port=port, user_id=user_id)
def dumps(self, user_id: typing.Optional[bytes] = None) -> bytes:
"""Packs the instance into a raw binary in the appropriate form.
Args:
user_id: Optional user ID as an override, if not provided the instance's
will be used, if none was provided at initialization an error is raised.
Returns:
The packed request.
Raises:
SOCKSError: If no user was specified in this call or on initialization.
"""
user_id = user_id or self.user_id
if user_id is None:
raise SOCKSError("SOCKS4 requires a user_id, none was specified")
return b"".join(
[
b"\x04",
self.command,
(self.port).to_bytes(2, byteorder="big"),
b"\x00\x00\x00\xFF", # arbitrary final non-zero byte
user_id,
b"\x00",
self.addr,
b"\x00",
]
)
class SOCKS4Reply(typing.NamedTuple):
"""Encapsulates a reply from the SOCKS4 proxy server
Args:
reply_code: The code representing the type of reply.
port: The port number returned.
addr: Optional IP address returned.
"""
reply_code: SOCKS4ReplyCode
port: int
addr: typing.Optional[str]
@classmethod
def loads(cls, data: bytes) -> "SOCKS4Reply":
"""Unpacks the reply data into an instance.
Returns:
The unpacked reply instance.
Raises:
ProtocolError: If the data does not match the spec.
"""
if len(data) != 8 or data[0:1] != b"\x00":
raise ProtocolError("Malformed reply")
try:
return cls(
reply_code=SOCKS4ReplyCode(data[1:2]),
port=int.from_bytes(data[2:4], byteorder="big"),
addr=decode_address(AddressType.IPV4, data[4:8]),
)
except ValueError as exc:
raise ProtocolError("Malformed reply") from exc
class SOCKS4Connection:
"""Encapsulates a SOCKS4 and SOCKS4A connection.
Packs request objects into data suitable to be send and unpacks reply
data into their appropriate reply objects.
Args:
user_id: The user ID to be sent as part of the requests.
"""
def __init__(self, user_id: bytes):
self.user_id = user_id
self._data_to_send = bytearray()
self._received_data = bytearray()
def send(self, request: typing.Union[SOCKS4Request, SOCKS4ARequest]) -> None:
"""Packs a request object and adds it to the send data buffer.
Args:
request: The request instance to be packed.
"""
user_id = request.user_id or self.user_id
self._data_to_send += request.dumps(user_id=user_id)
def receive_data(self, data: bytes) -> SOCKS4Reply:
"""Unpacks response data into a reply object.
Args:
data: The raw response data from the proxy server.
Returns:
The appropriate reply object.
"""
self._received_data += data
return SOCKS4Reply.loads(bytes(self._received_data))
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