path: root/client/3rd-party/client.py

# ai-generated, do not trust (i don't like python)

import asyncio
import aiohttp
import hashlib
from ecdsa import SigningKey, NIST256p
from typing import List, Tuple, Optional, Dict
from dataclasses import dataclass
import struct

MAX_BUFFER = 1 << 18
MAX_BODY = 1 << 22
MAX_RECIPIENTS = 32
PACKET_SIZE = 4 + 32 + (1 + 8 + 32 + 32 + 1 + (32 * MAX_RECIPIENTS) + 65536 + 64)
RELAY_NEIGHBOR_WIRE_LEN = 1 + 1 + 1 + 4 + 16 + 2

class MpClientError:
    SUCCESS = 0
    NETWORK_FAIL = 1
    CONNECTION_FAILED = 2
    SERVER_BUSY = 3
    INVALID_RESPONSE = 4
    DESERIALIZE_FAILED = 5
    STRING_ENCODE_FAILED = 6
    BAD_ARGUMENT = 7

@dataclass
class RelayNeighbor:
    active: bool
    has_ipv4: bool
    has_ipv6: bool
    v4_addr: bytes
    v6_addr: bytes
    port: int

@dataclass
class RelayInfo:
    major_version: int
    minor_version: int
    motd: bytes
    capabilities: Dict[str, any]
    statistics: Dict[str, any]

@dataclass
class DataPacketContents:
    encrypted: bool
    created_time_ms: int
    author: bytes
    author_signing_pub: bytes
    recipient_count: int
    recipients: List[bytes]
    data: bytes
    signature: bytes

@dataclass
class DataPacket:
    nonce: int
    hash: bytes
    contents: DataPacketContents

@dataclass
class DataPacketList:
    packets: List[DataPacket]
    count: int
    new_offset: int

@dataclass
class KeyPair:
    encryption_pub: bytes
    encryption_priv: bytes
    signing_pub: bytes
    signing_priv: bytes

class MoonPasteClient:
    def __init__(self):
        self.ip: Optional[str] = None
        self.port: Optional[int] = None

    def connect(self, ip: str, port: int) -> None:
        self.ip = ip
        self.port = port

    def _check_connected(self) -> bool:
        return self.ip is not None and self.port is not None

    def _read_u16_le(self, buffer: bytes, offset: int) -> int:
        return struct.unpack_from('<H', buffer, offset)[0]

    def _write_u16_le(self, buffer: bytearray, value: int, offset: int) -> None:
        struct.pack_into('<H', buffer, offset, value)

    def _write_u32_le(self, buffer: bytearray, value: int, offset: int) -> None:
        struct.pack_into('<I', buffer, offset, value)

    def _read_u32_le(self, buffer: bytes, offset: int) -> int:
        return struct.unpack_from('<I', buffer, offset)[0]

    def _read_u64_le(self, buffer: bytes, offset: int) -> int:
        return struct.unpack_from('<Q', buffer, offset)[0]

    def _write_u64_le(self, buffer: bytearray, value: int, offset: int) -> None:
        struct.pack_into('<Q', buffer, offset, value)

    def _check_leading_zeros(self, hash_bytes: bytes, required_zeros: int) -> bool:
        if required_zeros == 0:
            return True
        full_bytes = required_zeros // 8
        remaining_bits = required_zeros % 8
        for i in range(full_bytes):
            if hash_bytes[i] != 0:
                return False
        if remaining_bits > 0:
            mask = 0xFF << (8 - remaining_bits)
            if hash_bytes[full_bytes] & mask != 0:
                return False
        return True

    async def _sha256(self, data: bytes) -> bytes:
        return hashlib.sha256(data).digest()

    async def generate_key_pair(self) -> KeyPair:
        enc_key = SigningKey.generate(curve=NIST256p)
        sign_key = SigningKey.generate(curve=NIST256p)
        return KeyPair(
            encryption_pub=enc_key.verifying_key.to_string()[:32],
            encryption_priv=enc_key.to_string()[:32],
            signing_pub=sign_key.verifying_key.to_string()[:32],
            signing_priv=sign_key.to_string()[:32]
        )

    async def _sign_data(self, data: bytes, private_key: bytes) -> bytes:
        sk = SigningKey.from_string(private_key, curve=NIST256p)
        return sk.sign(data, hashfunc=hashlib.sha256)[:64]

    def _s2relay_neighbor(self, stream: bytes) -> Tuple[RelayNeighbor, bool]:
        if len(stream) != RELAY_NEIGHBOR_WIRE_LEN:
            return RelayNeighbor(False, False, False, b'\x00' * 4, b'\x00' * 16, 0), False
        return RelayNeighbor(
            active=stream[0] != 0,
            has_ipv4=stream[1] != 0,
            has_ipv6=stream[2] != 0,
            v4_addr=stream[3:7],
            v6_addr=stream[7:23],
            port=self._read_u16_le(stream, 23)
        ), True

    def _s2relay_info(self, stream: bytes) -> Tuple[RelayInfo, bool]:
        expected_len = 1 + 1 + 128 + 2 + 1 + 1 + 1 + 2 + 8 + (4 * RELAY_NEIGHBOR_WIRE_LEN)
        if len(stream) != expected_len:
            return RelayInfo(0, 0, b'\x00' * 128, {'retention_hours': 0, 'can_store': False, 'leading_zeros': 0, 'accepts_plain': False}, {'neighbors': 0, 'held_packets': 0, 'quick_connect': []}), False
        off = 0
        info = RelayInfo(
            major_version=stream[off],
            minor_version=stream[off + 1],
            motd=stream[off + 2:off + 130],
            capabilities={
                'retention_hours': self._read_u16_le(stream, off + 130),
                'can_store': stream[off + 132] != 0,
                'leading_zeros': stream[off + 133],
                'accepts_plain': stream[off + 134] != 0
            },
            statistics={
                'neighbors': self._read_u16_le(stream, off + 135),
                'held_packets': self._read_u64_le(stream, off + 137),
                'quick_connect': []
            }
        )
        off += 145
        for _ in range(4):
            neighbor, success = self._s2relay_neighbor(stream[off:off + RELAY_NEIGHBOR_WIRE_LEN])
            if not success:
                return RelayInfo(0, 0, b'\x00' * 128, {'retention_hours': 0, 'can_store': False, 'leading_zeros': 0, 'accepts_plain': False}, {'neighbors': 0, 'held_packets': 0, 'quick_connect': []}), False
            info.statistics['quick_connect'].append(neighbor)
            off += RELAY_NEIGHBOR_WIRE_LEN
        return info, True

    def _s2data_packet(self, stream: bytes) -> Tuple[DataPacket, bool]:
        if len(stream) != PACKET_SIZE:
            return DataPacket(0, b'\x00' * 32, DataPacketContents(False, 0, b'\x00' * 32, b'\x00' * 32, 0, [], b'\x00' * 65536, b'\x00' * 64)), False
        off = 0
        packet = DataPacket(
            nonce=self._read_u32_le(stream, off),
            hash=stream[off + 4:off + 36],
            contents=DataPacketContents(
                encrypted=stream[off + 36] != 0,
                created_time_ms=self._read_u64_le(stream, off + 37),
                author=stream[off + 45:off + 77],
                author_signing_pub=stream[off + 77:off + 109],
                recipient_count=stream[off + 109],
                recipients=[],
                data=stream[off + 110 + 32 * MAX_RECIPIENTS:off + 110 + 32 * MAX_RECIPIENTS + 65536],
                signature=stream[off + 110 + 32 * MAX_RECIPIENTS + 65536:off + 110 + 32 * MAX_RECIPIENTS + 65536 + 64]
            )
        )
        off += 110
        for _ in range(MAX_RECIPIENTS):
            packet.contents.recipients.append(stream[off:off + 32])
            off += 32
        return packet, True

    def _data_packet2s(self, data: DataPacket) -> Optional[bytes]:
        if data.contents.recipient_count > MAX_RECIPIENTS:
            return None
        stream = bytearray(PACKET_SIZE)
        off = 0
        self._write_u32_le(stream, data.nonce, off)
        off += 4
        stream[off:off + 32] = data.hash
        off += 32
        stream[off] = 1 if data.contents.encrypted else 0
        off += 1
        self._write_u64_le(stream, data.contents.created_time_ms, off)
        off += 8
        stream[off:off + 32] = data.contents.author
        off += 32
        stream[off:off + 32] = data.contents.author_signing_pub
        off += 32
        stream[off] = data.contents.recipient_count
        off += 1
        for i in range(MAX_RECIPIENTS):
            stream[off:off + 32] = data.contents.recipients[i] if i < len(data.contents.recipients) else b'\x00' * 32
            off += 32
        stream[off:off + 65536] = data.contents.data
        off += 65536
        stream[off:off + 64] = data.contents.signature
        return bytes(stream)

    async def _client_request(self, data: bytes, data_len: int) -> Tuple[Optional[bytes], int, int]:
        if not self._check_connected():
            return None, 0, MpClientError.CONNECTION_FAILED
        is_ipv6 = ':' in self.ip
        url = f"http://[{self.ip}]:{self.port}/" if is_ipv6 else f"http://{self.ip}:{self.port}/"
        async with aiohttp.ClientSession() as session:
            try:
                async with session.post(url, data=data[:data_len], headers={'Content-Length': str(data_len), 'Connection': 'close'}) as response:
                    if response.status == 418:
                        return None, 0, MpClientError.SERVER_BUSY
                    if response.status != 200:
                        return None, 0, MpClientError.NETWORK_FAIL
                    response_data = await response.read()
                    content_length = int(response.headers.get('Content-Length', '0'))
                    if content_length > 0 and len(response_data) != content_length:
                        return None, 0, MpClientError.RECV_FAILED
                    if len(response_data) > MAX_BODY:
                        return None, 0, MpClientError.INVALID_RESPONSE
                    return response_data, len(response_data), MpClientError.SUCCESS
            except aiohttp.ClientError:
                return None, 0, MpClientError.CONNECTION_FAILED

    async def get_server_info(self) -> Tuple[RelayInfo, int]:
        data = b'\x10'
        response, resp_len, error = await self._client_request(data, len(data))
        if error != MpClientError.SUCCESS:
            return RelayInfo(0, 0, b'\x00' * 128, {'retention_hours': 0, 'can_store': False, 'leading_zeros': 0, 'accepts_plain': False}, {'neighbors': 0, 'held_packets': 0, 'quick_connect': []}), error
        if resp_len < 1 or response[0] != 0x17:
            return RelayInfo(0, 0, b'\x00' * 128, {'retention_hours': 0, 'can_store': False, 'leading_zeros': 0, 'accepts_plain': False}, {'neighbors': 0, 'held_packets': 0, 'quick_connect': []}), MpClientError.INVALID_RESPONSE
        info, success = self._s2relay_info(response[1:resp_len])
        return info, MpClientError.SUCCESS if success else MpClientError.DESERIALIZE_FAILED

    async def get_server_neighbors(self, count: int) -> Tuple[List[RelayNeighbor], int, int]:
        data = bytearray(2)
        data[0] = 0x11
        data[1] = count
        response, resp_len, error = await self._client_request(data, len(data))
        if error != MpClientError.SUCCESS:
            return [], 0, error
        if resp_len < 1 or response[0] != 0x18:
            return [], 0, MpClientError.INVALID_RESPONSE
        num_neighbors = (resp_len - 1) // RELAY_NEIGHBOR_WIRE_LEN
        neighbors = []
        actual_count = 0
        for i in range(num_neighbors):
            neighbor, success = self._s2relay_neighbor(response[1 + i * RELAY_NEIGHBOR_WIRE_LEN:1 + (i + 1) * RELAY_NEIGHBOR_WIRE_LEN])
            if success:
                neighbors.append(neighbor)
                actual_count += 1
        return neighbors, actual_count, MpClientError.SUCCESS

    async def push_packet(self, packet: DataPacket) -> Tuple[bool, int]:
        send_data = self._data_packet2s(packet)
        if send_data is None:
            return False, MpClientError.BAD_ARGUMENT
        full_data = bytearray(1 + PACKET_SIZE)
        full_data[0] = 0x13
        full_data[1:] = send_data
        response, resp_len, error = await self._client_request(full_data, len(full_data))
        if error != MpClientError.SUCCESS:
            return False, error
        if resp_len == 2 and response[0] == 0x1a:
            return response[1] != 0, MpClientError.SUCCESS
        return False, MpClientError.INVALID_RESPONSE

    async def get_packet(self, packet_hash: bytes) -> Tuple[DataPacket, int]:
        send_data = bytearray(1 + 32)
        send_data[0] = 0x14
        send_data[1:33] = packet_hash
        response, resp_len, error = await self._client_request(send_data, len(send_data))
        if error != MpClientError.SUCCESS:
            return DataPacket(0, b'\x00' * 32, DataPacketContents(False, 0, b'\x00' * 32, b'\x00' * 32, 0, [], b'\x00' * 65536, b'\x00' * 64)), error
        if resp_len < 1 or response[0] != 0x1b:
            return DataPacket(0, b'\x00' * 32, DataPacketContents(False, 0, b'\x00' * 32, b'\x00' * 32, 0, [], b'\x00' * 65536, b'\x00' * 64)), MpClientError.INVALID_RESPONSE
        packet, success = self._s2data_packet(response[1:resp_len])
        return packet, MpClientError.SUCCESS if success else MpClientError.DESERIALIZE_FAILED

    async def get_firehose_packets(self, offset: int, count: int) -> Tuple[DataPacketList, int]:
        send_data = bytearray(1 + 1 + 8)
        send_data[0] = 0x12
        send_data[1] = count
        self._write_u64_le(send_data, offset, 2)
        response, resp_len, error = await self._client_request(send_data, len(send_data))
        if error != MpClientError.SUCCESS:
            return DataPacketList([], 0, offset), error
        if resp_len < 9 or response[0] != 0x19:
            return DataPacketList([], 0, offset), MpClientError.INVALID_RESPONSE
        new_offset = self._read_u64_le(response, 1)
        available_bytes = resp_len - 9
        available_packets = available_bytes // PACKET_SIZE
        packets = []
        packet_count = 0
        for i in range(available_packets):
            packet, success = self._s2data_packet(response[9 + i * PACKET_SIZE:9 + (i + 1) * PACKET_SIZE])
            if success:
                packets.append(packet)
                packet_count += 1
        return DataPacketList(packets, packet_count, new_offset), MpClientError.SUCCESS

    async def create_packet_from_bytes(
        self,
        input_bytes: bytes,
        input_len: int,
        recipient_count: int,
        recipients: List[bytes],
        author_keys: KeyPair
    ) -> DataPacket:
        packet = DataPacket(
            nonce=0,
            hash=b'\x00' * 32,
            contents=DataPacketContents(
                encrypted=False,
                created_time_ms=int.from_bytes(struct.pack('<Q', int.from_bytes(b'\x00' * 8, 'little') + int(1000 * (await asyncio.get_event_loop().time()))), 'little'),
                author=author_keys.encryption_pub,
                author_signing_pub=author_keys.signing_pub,
                recipient_count=min(recipient_count, MAX_RECIPIENTS),
                recipients=[b'\x00' * 32 for _ in range(MAX_RECIPIENTS)],
                data=b'\x00' * 65536,
                signature=b'\x00' * 64
            )
        )
        copy_len = min(input_len, 65536)
        if input_bytes and copy_len > 0:
            packet.contents.data = input_bytes[:copy_len] + b'\x00' * (65536 - copy_len)
        if recipients:
            for i in range(packet.contents.recipient_count):
                packet.contents.recipients[i] = recipients[i]
        msg = bytearray(1 + 8 + 32 + 32 + 1 + 32 * MAX_RECIPIENTS + 65536)
        msg[0] = 1 if packet.contents.encrypted else 0
        self._write_u64_le(msg, packet.contents.created_time_ms, 1)
        msg[9:41] = packet.contents.author
        msg[41:73] = packet.contents.author_signing_pub
        msg[73] = packet.contents.recipient_count
        for i in range(MAX_RECIPIENTS):
            msg[74 + 32 * i:74 + 32 * (i + 1)] = packet.contents.recipients[i]
        msg[74 + 32 * MAX_RECIPIENTS:74 + 32 * MAX_RECIPIENTS + 65536] = packet.contents.data
        packet.contents.signature = await self._sign_data(msg, author_keys.signing_priv)
        return packet

    async def create_packet_from_string(
        self,
        input_str: str,
        recipient_count: int,
        recipients: List[bytes],
        author_keys: KeyPair
    ) -> Tuple[DataPacket, int]:
        try:
            input_bytes = input_str.encode('utf-8')
            packet = await self.create_packet_from_bytes(input_bytes, len(input_bytes), recipient_count, recipients, author_keys)
            return packet, MpClientError.SUCCESS
        except UnicodeEncodeError:
            return DataPacket(0, b'\x00' * 32, DataPacketContents(False, 0, b'\x00' * 32, b'\x00' * 32, 0, [], b'\x00' * 65536, b'\x00' * 64)), MpClientError.STRING_ENCODE_FAILED

    async def run_pow(self, packet: DataPacket, leading_zeros: int) -> DataPacket:
        hash_input = bytearray(4 + PACKET_SIZE - 4)
        contents_buffer = hash_input[4:]
        contents_buffer[0] = 1 if packet.contents.encrypted else 0
        self._write_u64_le(contents_buffer, packet.contents.created_time_ms, 1)
        contents_buffer[9:41] = packet.contents.author
        contents_buffer[41:73] = packet.contents.author_signing_pub
        contents_buffer[73] = packet.contents.recipient_count
        for i in range(MAX_RECIPIENTS):
            contents_buffer[74 + 32 * i:74 + 32 * (i + 1)] = packet.contents.recipients[i]
        contents_buffer[74 + 32 * MAX_RECIPIENTS:74 + 32 * MAX_RECIPIENTS + 65536] = packet.contents.data
        contents_buffer[74 + 32 * MAX_RECIPIENTS + 65536:74 + 32 * MAX_RECIPIENTS + 65536 + 64] = packet.contents.signature
        nonce = 0
        while True:
            self._write_u32_le(hash_input, nonce, 0)
            hash_result = await self._sha256(hash_input)
            if self._check_leading_zeros(hash_result, leading_zeros):
                packet.nonce = nonce
                packet.hash = hash_result
                break
            nonce += 1
        return packet

    async def query_author_packets(
        self,
        author_keys: KeyPair,
        count: int,
        offset: int
    ) -> Tuple[DataPacketList, int]:
        msg = bytearray(32 + 32 + 2 + 8)
        msg[0:32] = author_keys.encryption_pub
        msg[32:64] = author_keys.signing_pub
        self._write_u16_le(msg, count, 64)
        self._write_u64_le(msg, offset, 66)
        sig = await self._sign_data(msg, author_keys.signing_priv)
        send_data = bytearray(1 + 32 + 32 + 64 + 2 + 8)
        off = 0
        send_data[off] = 0x16
        off += 1
        send_data[off:off + 32] = author_keys.encryption_pub
        off += 32
        send_data[off:off + 32] = author_keys.signing_pub
        off += 32
        send_data[off:off + 64] = sig
        off += 64
        self._write_u16_le(send_data, count, off)
        off += 2
        self._write_u64_le(send_data, offset, off)
        response, resp_len, error = await self._client_request(send_data, len(send_data))
        if error != MpClientError.SUCCESS:
            return DataPacketList([], 0, offset), error
        if resp_len < 11 or response[0] != 0x1d:
            return DataPacketList([], 0, offset), MpClientError.INVALID_RESPONSE
        resp_count = self._read_u16_le(response, 1)
        new_offset = self._read_u64_le(response, 3)
        expected_len = 11 + resp_count * PACKET_SIZE
        if resp_len != expected_len:
            return DataPacketList([], 0, offset), MpClientError.INVALID_RESPONSE
        packets = []
        packet_count = 0
        for i in range(resp_count):
            packet, success = self._s2data_packet(response[11 + i * PACKET_SIZE:11 + (i + 1) * PACKET_SIZE])
            if success:
                packets.append(packet)
                packet_count += 1
        return DataPacketList(packets, packet_count, new_offset), MpClientError.SUCCESS