nix-ip-utils/lib/cidr.nix

{
  lib,
  internal,
  ip,
}:

let
  inherit (builtins)
    elemAt
    filter
    map
    split
    ;
  inherit (lib) toInt;
  inherit (internal) pow2 intRange;

in
rec {
  /**
    Parse a CIDR string into { address; prefixLen }.
    Also accepts attrsets (returns as-is).

    # Type

    ```
    parse :: (String | { address :: Int; prefixLen :: Int; }) -> { address :: Int; prefixLen :: Int; }
    ```

    # Example

    ```nix
    parse "192.168.1.0/24"
    => { address = 3232235776; prefixLen = 24; }
    ```
  */
  parse =
    cidr:
    if builtins.isAttrs cidr then
      cidr
    else
      let
        parts = filter builtins.isString (split "/" cidr);
        addr = ip.parse (elemAt parts 0);
        prefix = toInt (elemAt parts 1);
      in
      {
        address = addr;
        prefixLen = prefix;
      };

  /**
    Create a CIDR from IP and prefix length.
    Returns { address; prefixLen }.

    # Type

    ```
    make :: (String | Int) -> Int -> { address :: Int; prefixLen :: Int; }
    ```

    # Example

    ```nix
    make "192.168.1.0" 24
    => { address = 3232235776; prefixLen = 24; }
    ```
  */
  make = addr: prefixLen: {
    address = ip.parse addr;
    inherit prefixLen;
  };

  /**
    Format a CIDR as string.

    # Type

    ```
    format :: (String | { address :: Int; prefixLen :: Int; }) -> String
    ```

    # Example

    ```nix
    format { address = 3232235776; prefixLen = 24; }
    => "192.168.1.0/24"
    ```
  */
  format =
    cidr:
    let
      c = parse cidr;
    in
    "${ip.format c.address}/${toString c.prefixLen}";

  /**
    Calculate netmask from prefix length (as integer).

    # Type

    ```
    netmask :: Int -> Int
    ```

    # Example

    ```nix
    netmask 24
    => 4294967040
    ```
  */
  netmask =
    prefixLen:
    if prefixLen == 0 then
      0
    else
      let
        # Create mask with prefixLen 1s followed by (32-prefixLen) 0s
        hostBits = 32 - prefixLen;
      in
      4294967295 - (pow2 hostBits - 1);

  /**
    Calculate netmask from prefix length (as string).

    # Type

    ```
    netmaskStr :: Int -> String
    ```

    # Example

    ```nix
    netmaskStr 24
    => "255.255.255.0"
    ```
  */
  netmaskStr = prefixLen: ip.format (netmask prefixLen);

  /**
    Calculate wildcard mask (inverse of netmask).

    # Type

    ```
    wildcard :: Int -> Int
    ```

    # Example

    ```nix
    wildcard 24
    => 255
    ```
  */
  wildcard =
    prefixLen:
    let
      hostBits = 32 - prefixLen;
    in
    pow2 hostBits - 1;

  /**
    Calculate wildcard mask as string.

    # Type

    ```
    wildcardStr :: Int -> String
    ```

    # Example

    ```nix
    wildcardStr 24
    => "0.0.0.255"
    ```
  */
  wildcardStr = prefixLen: ip.format (wildcard prefixLen);

  /**
    Get network address of a CIDR (as integer).

    # Type

    ```
    network :: (String | { address :: Int; prefixLen :: Int; }) -> Int
    ```

    # Example

    ```nix
    network "192.168.1.50/24"
    => 3232235776
    ```
  */
  network =
    cidr:
    let
      c = parse cidr;
      mask = netmask c.prefixLen;
    in
    builtins.bitAnd c.address mask;

  /**
    Get network address of a CIDR (as string).

    # Type

    ```
    networkStr :: (String | { address :: Int; prefixLen :: Int; }) -> String
    ```

    # Example

    ```nix
    networkStr "192.168.1.50/24"
    => "192.168.1.0"
    ```
  */
  networkStr = cidr: ip.format (network cidr);

  /**
    Get broadcast address of a CIDR (as integer).

    # Type

    ```
    broadcast :: (String | { address :: Int; prefixLen :: Int; }) -> Int
    ```

    # Example

    ```nix
    broadcast "192.168.1.0/24"
    => 3232236031
    ```
  */
  broadcast =
    cidr:
    let
      c = parse cidr;
      wild = wildcard c.prefixLen;
      net = network cidr;
    in
    builtins.bitOr net wild;

  /**
    Get broadcast address of a CIDR (as string).

    # Type

    ```
    broadcastStr :: (String | { address :: Int; prefixLen :: Int; }) -> String
    ```

    # Example

    ```nix
    broadcastStr "192.168.1.0/24"
    => "192.168.1.255"
    ```
  */
  broadcastStr = cidr: ip.format (broadcast cidr);

  /**
    Get first usable host address (as integer).
    For /31 and /32, returns network address.

    # Type

    ```
    firstHost :: (String | { address :: Int; prefixLen :: Int; }) -> Int
    ```

    # Example

    ```nix
    firstHost "192.168.1.0/24"
    => 3232235777
    ```
  */
  firstHost =
    cidr:
    let
      c = parse cidr;
      net = network cidr;
    in
    if c.prefixLen >= 31 then net else net + 1;

  /**
    Get first usable host address (as string).

    # Type

    ```
    firstHostStr :: (String | { address :: Int; prefixLen :: Int; }) -> String
    ```

    # Example

    ```nix
    firstHostStr "192.168.1.0/24"
    => "192.168.1.1"
    ```
  */
  firstHostStr = cidr: ip.format (firstHost cidr);

  /**
    Get last usable host address (as integer).
    For /31 and /32, returns broadcast address.

    # Type

    ```
    lastHost :: (String | { address :: Int; prefixLen :: Int; }) -> Int
    ```

    # Example

    ```nix
    lastHost "192.168.1.0/24"
    => 3232236030
    ```
  */
  lastHost =
    cidr:
    let
      c = parse cidr;
      bcast = broadcast cidr;
    in
    if c.prefixLen >= 31 then bcast else bcast - 1;

  /**
    Get last usable host address (as string).

    # Type

    ```
    lastHostStr :: (String | { address :: Int; prefixLen :: Int; }) -> String
    ```

    # Example

    ```nix
    lastHostStr "192.168.1.0/24"
    => "192.168.1.254"
    ```
  */
  lastHostStr = cidr: ip.format (lastHost cidr);

  /**
    Get total number of addresses in CIDR.

    # Type

    ```
    size :: (String | { address :: Int; prefixLen :: Int; }) -> Int
    ```

    # Example

    ```nix
    size "192.168.1.0/24"
    => 256
    ```
  */
  size =
    cidr:
    let
      c = parse cidr;
    in
    pow2 (32 - c.prefixLen);

  /**
    Get number of usable host addresses.
    Excludes network and broadcast for normal subnets.
    /31 returns 2, /32 returns 1 (special cases per RFC 3021).

    # Type

    ```
    hostCount :: (String | { address :: Int; prefixLen :: Int; }) -> Int
    ```

    # Example

    ```nix
    hostCount "192.168.1.0/24"
    => 254
    ```
  */
  hostCount =
    cidr:
    let
      c = parse cidr;
      total = size cidr;
    in
    if c.prefixLen == 32 then
      1
    else if c.prefixLen == 31 then
      2
    else if total <= 2 then
      0
    else
      total - 2;

  /**
    Check if an IP address is within a CIDR.

    # Type

    ```
    contains :: (String | { address :: Int; prefixLen :: Int; }) -> (String | Int) -> Bool
    ```

    # Example

    ```nix
    contains "10.0.0.0/8" "10.1.2.3"
    => true
    ```
  */
  contains =
    cidr: addr:
    let
      ipInt = ip.parse addr;
      net = network cidr;
      bcast = broadcast cidr;
    in
    ipInt >= net && ipInt <= bcast;

  /**
    Check if cidr2 is a subnet of cidr1.

    # Type

    ```
    isSubnetOf :: (String | { address :: Int; prefixLen :: Int; }) -> (String | { address :: Int; prefixLen :: Int; }) -> Bool
    ```

    # Example

    ```nix
    isSubnetOf "10.0.0.0/8" "10.1.0.0/16"
    => true
    ```
  */
  isSubnetOf =
    cidr1: cidr2:
    let
      c1 = parse cidr1;
      c2 = parse cidr2;
      net1 = network cidr1;
      bcast1 = broadcast cidr1;
      net2 = network cidr2;
      bcast2 = broadcast cidr2;
    in
    c2.prefixLen >= c1.prefixLen && net2 >= net1 && bcast2 <= bcast1;

  /**
    Get the Nth host in a CIDR (0-indexed from network address).

    # Type

    ```
    host :: (String | { address :: Int; prefixLen :: Int; }) -> Int -> Int
    ```

    # Example

    ```nix
    host "192.168.1.0/24" 5
    => 3232235781
    ```
  */
  host = cidr: n: (network cidr) + n;

  /**
    Get the Nth host in a CIDR (as string).

    # Type

    ```
    hostStr :: (String | { address :: Int; prefixLen :: Int; }) -> Int -> String
    ```

    # Example

    ```nix
    hostStr "192.168.1.0/24" 5
    => "192.168.1.5"
    ```
  */
  hostStr = cidr: n: ip.format (host cidr n);

  /**
    Subdivide a CIDR into smaller subnets.
    Returns the Nth subnet with the new prefix length.

    # Type

    ```
    subnet :: (String | { address :: Int; prefixLen :: Int; }) -> Int -> Int -> { address :: Int; prefixLen :: Int; }
    ```

    # Example

    ```nix
    subnet "192.168.0.0/24" 26 0
    => { address = 3232235520; prefixLen = 26; }
    ```
  */
  subnet =
    cidr: newPrefix: n:
    let
      net = network cidr;
      subnetSize = pow2 (32 - newPrefix);
      subnetStart = net + (n * subnetSize);
    in
    {
      address = subnetStart;
      prefixLen = newPrefix;
    };

  /**
    Get all subnets when subdividing.

    # Type

    ```
    subnets :: (String | { address :: Int; prefixLen :: Int; }) -> Int -> [{ address :: Int; prefixLen :: Int; }]
    ```

    # Example

    ```nix
    subnets "192.168.0.0/24" 26
    => [ { address = 3232235520; prefixLen = 26; } ... ] # 4 subnets
    ```
  */
  subnets =
    cidr: newPrefix:
    let
      c = parse cidr;
      numSubnets = pow2 (newPrefix - c.prefixLen);
    in
    map (n: subnet cidr newPrefix n) (intRange 0 (numSubnets - 1));

  /**
    Get all subnets as strings.

    # Type

    ```
    subnetsStr :: (String | { address :: Int; prefixLen :: Int; }) -> Int -> [String]
    ```

    # Example

    ```nix
    subnetsStr "192.168.0.0/24" 26
    => [ "192.168.0.0/26" "192.168.0.64/26" "192.168.0.128/26" "192.168.0.192/26" ]
    ```
  */
  subnetsStr = cidr: newPrefix: map format (subnets cidr newPrefix);

  /**
    Check if two CIDRs overlap.

    # Type

    ```
    overlaps :: (String | { address :: Int; prefixLen :: Int; }) -> (String | { address :: Int; prefixLen :: Int; }) -> Bool
    ```

    # Example

    ```nix
    overlaps "192.168.1.0/24" "192.168.1.128/25"
    => true
    ```
  */
  overlaps =
    cidr1: cidr2:
    let
      net1 = network cidr1;
      bcast1 = broadcast cidr1;
      net2 = network cidr2;
      bcast2 = broadcast cidr2;
    in
    !(bcast1 < net2 || bcast2 < net1);

  /**
    Get prefix length from number of required hosts.

    # Type

    ```
    prefixForHosts :: Int -> Int
    ```

    # Example

    ```nix
    prefixForHosts 100
    => 25
    ```
  */
  prefixForHosts =
    hosts:
    let
      findPrefix =
        prefix:
        if prefix < 0 then
          0
        else if
          hostCount {
            address = 0;
            prefixLen = prefix;
          } >= hosts
        then
          prefix
        else
          findPrefix (prefix - 1);
    in
    findPrefix 32;

  /**
    Get prefix length from number of required addresses.

    # Type

    ```
    prefixForSize :: Int -> Int
    ```

    # Example

    ```nix
    prefixForSize 256
    => 24
    ```
  */
  prefixForSize =
    addresses:
    let
      findPrefix =
        prefix:
        if prefix < 0 then
          0
        else if
          size {
            address = 0;
            prefixLen = prefix;
          } >= addresses
        then
          prefix
        else
          findPrefix (prefix - 1);
    in
    findPrefix 32;
}