#!/usr/bin/env python3
"""Calcula a mano el checksum de un datagrama UDP sobre IPv4, palabra por palabra.

    python3 udp_checksum.py <ip_origen> <ip_destino> <puerto_origen> <puerto_destino> <datos>

Sigue el RFC 768: suma de complemento a uno del pseudo-encabezado, el encabezado UDP
(con el checksum en cero) y los datos, y al final el complemento. Si el resultado da
cero, se escribe 0xffff, porque en UDP sobre IPv4 un cero quiere decir "sin checksum".
"""
import argparse
import ipaddress
import struct

PROTO_UDP = 17


def pseudo_header(source, destination, udp_length):
    return (ipaddress.IPv4Address(source).packed + ipaddress.IPv4Address(destination).packed
            + struct.pack("!BBH", 0, PROTO_UDP, udp_length))


def words(data):
    """Parte en palabras de 16 bits. Si el largo es impar, completa con un byte en cero."""
    if len(data) % 2:
        data += b"\0"
    return [int.from_bytes(data[i:i + 2], "big") for i in range(0, len(data), 2)]


def ones_complement_sum(values):
    total = sum(values)
    while total > 0xffff:
        total = (total & 0xffff) + (total >> 16)
    return total


def udp_checksum(source, destination, source_port, destination_port, payload):
    length = 8 + len(payload)
    data = (pseudo_header(source, destination, length)
            + struct.pack("!HHHH", source_port, destination_port, length, 0) + payload)
    result = ~ones_complement_sum(words(data)) & 0xffff
    return result if result != 0 else 0xffff


def explain(source, destination, source_port, destination_port, payload):
    """Devuelve la cuenta como texto, en líneas."""
    length = 8 + len(payload)
    parts = [("pseudo-encabezado", pseudo_header(source, destination, length)),
             ("encabezado UDP", struct.pack("!HHHH", source_port, destination_port, length, 0)),
             ("datos", payload)]
    lines = []
    all_words = []
    for name, data in parts:
        block = words(data)
        all_words += block
        lines.append(f"{name:<18} " + " ".join(f"{word:04x}" for word in block))
    total = sum(all_words)
    lines.append(f"{'suma':<18} {total:x}")
    folded = ones_complement_sum(all_words)
    if total != folded:
        lines.append(f"{'con acarreos':<18} {folded:04x}")
    lines.append(f"{'complemento':<18} {~folded & 0xffff:04x}")
    lines.append(f"{'checksum':<18} "
                 f"{udp_checksum(source, destination, source_port, destination_port, payload):04x}")
    return lines


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("source")
    parser.add_argument("destination")
    parser.add_argument("source_port", type=int)
    parser.add_argument("destination_port", type=int)
    parser.add_argument("payload")
    args = parser.parse_args()
    print("\n".join(explain(args.source, args.destination, args.source_port,
                            args.destination_port, args.payload.encode())))


if __name__ == "__main__":
    main()