Address calculator for netfilter/ip6t_NPT

nptcalc.py

import sys
import socket
import ipaddress

def csum_add(csum, addend):
    csum += addend
    if csum > 0xFFFF:
        return (csum&0xFFFF) + 1
    return csum & 0xFFFF

def csum_sub(csum, addend):
    return csum_add(csum, (~addend)&0xFFFF)

def bytes2word(data):
    if len(data) != 2:
        raise ValueError(f"len(data) (={len(data)}) must be 2")
    return (data[0] << 8) + data[1]

def word2bytes(word):
    result = bytearray(2)
    result[0] = (word >> 8) & 0xFF
    result[1] = word & 0xFF
    return bytes(result)

def appendword(word, suffix, prefixlen):
    if prefixlen < 0 or prefixlen > 16:
        raise ValueError(f"prefixlen (={prefixlen} must in 0~16)")
    suffixlen = 16-prefixlen
    suffixmask = (1<<suffixlen)-1
    prefixmask = ~suffixmask
    return (word&prefixmask)|(suffix&suffixmask)

def bnpt(bsrc, bdst, prefixlen):
    if len(bsrc) != 16:
        raise ValueError(f"len(bsrc) (={len(bsrc)}) must be 16")
    if len(bdst) != 16:
        raise ValueError(f"len(bdst) (={len(bdst)}) must be 16")
    translen = (prefixlen+15)>>4<<4
    if translen < 48:
        translen = 48

    bres = bytearray(16)
    csrc = 0
    cdst = 0
    cres = 0

    for i in range(0, 16, 2):
        wsrc = bytes2word(bsrc[i:i+2])
        wdst = bytes2word(bdst[i:i+2])

        csrc = csum_add(csrc, wsrc)

        if i < prefixlen//16*2:
            wres = wdst
        elif i == prefixlen//16*2:
            wres = appendword(wdst, wsrc, prefixlen-i*8)
        else:
            wres = wsrc

        cdst = csum_add(cdst, wres)

        if i == translen//16*2:
            if wsrc == 0xFFFF:
                translen += 16
                if translen >= 128:
                    raise ValueError("NPT does not support this src address")
            else:
                adjustment = csum_sub(csrc, cdst)
                wres = csum_add(wres, adjustment)

        bres[i:i+2] = word2bytes(wres)
        cres = csum_add(cres, wres)

    if csrc != cres:
        raise ValueError(f"Error: csrc (={csrc}) != cres (={cres})")

    return bres

def print_baddr(baddr):
    print(socket.inet_ntop(socket.AF_INET6, baddr))

def npt(src, dst):
    '''
    src: IPv6Address
    dst: IPv6Network
    '''
    
    bsrc = socket.inet_pton(socket.AF_INET6, str(src))
    bdst = socket.inet_pton(socket.AF_INET6, str(dst.network_address))
    prefixlen = dst.prefixlen
    bres = bnpt(bsrc, bdst, prefixlen)
    return ipaddress.IPv6Address(socket.inet_ntop(socket.AF_INET6, bres))


def main():
    if "-h" in sys.argv or "--help" in sys.argv:
        sys.exit(f"Usage: {sys.argv[0]} from_addr to_pfx")
        return
    if len(sys.argv) != 3:
        sys.exit(f"Error: expected 2 arguments, got {len(sys.argv)}")
    src = ipaddress.IPv6Address(sys.argv[1])
    dst = ipaddress.IPv6Network(sys.argv[2])
    print(npt(src, dst))

if __name__ == "__main__":
    main()