#!/usr/bin/env python3

"""
Test of offline signing using KSR and SKR with pre-planned KSK rollover and automatic ZSK rollover.
"""

import random
from subprocess import check_output, DEVNULL

from dnstest.utils import *
from dnstest.keys import Keymgr
from dnstest.test import Test
from dnstest.params import bind_dnssec_ksr

def cripple_skr(skr_in, skr_out):
    rrsigs_total = 9
    after_rrsig = -1000
    rrsig_now = 0
    rrsig_chosen = random.randint(1, rrsigs_total)
    with open(skr_in, "r") as fin:
        with open(skr_out, "w") as fout:
            for linein in fin:
                lineout = linein
                linesplit = linein.split()
                if len(linesplit) > 2 and linesplit[2] == "RRSIG":
                    after_rrsig = 0
                    rrsig_now += 1
                elif len(linesplit) > 3 and linesplit[3] == "RRSIG":
                    rrsig_now += 1
                    if rrsig_now == rrsig_chosen:
                        lineout = linein.lower() # Bind9-style RRSIG is crippled directly
                else:
                    after_rrsig += 1
                    if after_rrsig == 3 and rrsig_now == rrsig_chosen:
                        lineout = linein.lower() # this crippels the rrsig
                fout.write(lineout)

# check zone if keys are present and used for signing
def check_zone(server, zone, dnskeys, dnskey_rrsigs, soa_rrsigs, msg):
    qdnskeys = server.dig("example.com", "DNSKEY", bufsize=4096)
    found_dnskeys = qdnskeys.count("DNSKEY")

    qdnskeyrrsig = server.dig("example.com", "DNSKEY", dnssec=True, bufsize=4096)
    found_rrsigs = qdnskeyrrsig.count("RRSIG")

    qsoa = server.dig("example.com", "SOA", dnssec=True, bufsize=4096)
    found_soa_rrsigs = qsoa.count("RRSIG")

    check_log("DNSKEYs: %d (expected %d)" % (found_dnskeys, dnskeys))
    check_log("RRSIGs: %d (expected %d)" % (found_soa_rrsigs, soa_rrsigs))
    check_log("DNSKEY-RRSIGs: %d (expected %d)" % (found_rrsigs, dnskey_rrsigs))

    if found_dnskeys != dnskeys:
        set_err("BAD DNSKEY COUNT: " + msg)
        detail_log("!DNSKEYs not published and activated as expected: " + msg)

    if found_soa_rrsigs != soa_rrsigs:
        set_err("BAD RRSIG COUNT: " + msg)
        detail_log("!RRSIGs not published and activated as expected: " + msg)

    if found_rrsigs != dnskey_rrsigs:
        set_err("BAD DNSKEY RRSIG COUNT: " + msg)
        detail_log("!RRSIGs not published and activated as expected: " + msg)

    detail_log(SEP)

    server.zone_backup(zone, flush=True)
    server.zone_verify(zone)

def wait_for_rrsig_count(t, server, rrtype, rrsig_count, timeout):
    endtime = time.monotonic() + timeout - 0.5
    first = True
    while True:
        qdnskeyrrsig = server.dig("example.com", rrtype, dnssec=True, bufsize=4096)
        found_rrsigs = qdnskeyrrsig.count("RRSIG")
        if found_rrsigs == rrsig_count:
            break

        if first:
            first = False
            # Verify the zone instead of a dumb sleep
            server.zone_backup(zone, flush=True)
            server.zone_verify(zone)
        else:
            t.sleep(0.5)

        if time.monotonic() > endtime:
            break

def wait_for_dnskey_count(t, server, dnskey_count, timeout):
    endtime = time.monotonic() + timeout - 0.5
    first = True
    while True:
        qdnskeyrrsig = server.dig("example.com", "DNSKEY", dnssec=True, bufsize=4096)
        found_dnskeys = qdnskeyrrsig.count("DNSKEY")
        if found_dnskeys == dnskey_count:
            break

        if first:
            first = False
            # Verify the zone instead of a dumb sleep
            server.zone_backup(zone, flush=True)
            server.zone_verify(zone)
        else:
            t.sleep(0.5)

        if time.monotonic() > endtime:
            break

def zone_update(master, slave, zone, upd_master):
    server = master if upd_master else slave
    server.ctl("-f zone-flush", wait=True) # re-flush zone file so that it's suitable for generating radnom DDNS
    server.random_ddns(zone[0], allow_empty=True)

def writef(filename, contents):
    with open(filename, "w") as f:
        f.write(contents)

ON_SLAVE = random.choice([True, False])
IXFR = random.choice([True, False]) if ON_SLAVE else False
SIGNER_BIND = random.choice([True, False])

DNSSEC_KSR_VER = None if bind_dnssec_ksr is None or bind_dnssec_ksr == "" else check_output([ bind_dnssec_ksr, "-V" ], stderr=DEVNULL).decode('utf-8').split('.')[1]
if DNSSEC_KSR_VER is None or int(DNSSEC_KSR_VER) < 21:
    SIGNER_BIND = False

check_log("On-slave signing %s, IXFR enabled %s, Bind9 as signer %s, dnssec-ksr version %s" % (ON_SLAVE, IXFR, SIGNER_BIND, str(DNSSEC_KSR_VER)))

t = Test()

knot = t.server("knot")
ZONE = "example.com."
FUTURE = 55
TICK = 5
STARTUP = 10
NONSENSE = 4396

zone = t.zone(ZONE)
if ON_SLAVE:
    master = t.server("knot")
    t.link(zone, master, knot)
    if not IXFR:
        master.conf_zone(zone).journal_content = "none"
else:
    master = t.server("dummy")
    t.link(zone, knot)

knot.conf_zone(zone).zonefile_sync = 24 * 60 * 60

# ZSK side
knot.dnssec(zone).enable = True
knot.dnssec(zone).manual = True
knot.dnssec(zone).offline_ksk = True
knot.dnssec(zone).dnskey_ttl = 2
knot.dnssec(zone).zone_max_ttl = 3
knot.dnssec(zone).keytag_modulo = "0/2"
# optional
knot.dnssec(zone).algorithm = "ECDSAP384SHA384"
knot.dnssec(zone).zsk_lifetime = STARTUP + 6 * TICK # see ksk1 lifetime
knot.dnssec(zone).propagation_delay = TICK - 2
knot.dnssec(zone).cds_cdnskey_publish = "rollover"
knot.dnssec(zone).rrsig_lifetime = 15
knot.dnssec(zone).rrsig_refresh = 6
knot.dnssec(zone).rrsig_pre_refresh = 1
# options without any effect
knot.dnssec(zone).ksk_lifetime = NONSENSE

signer = t.server("bind" if SIGNER_BIND else "knot")
t.link(zone, signer)

# KSK side
signer.dnssec(zone).enable = True
signer.dnssec(zone).manual = True
signer.dnssec(zone).offline_ksk = True
signer.dnssec(zone).rrsig_refresh = 2
signer.dnssec(zone).keytag_modulo = "1/2"
# optional
signer.dnssec(zone).algorithm = "ECDSAP384SHA384"
signer.dnssec(zone).cds_cdnskey_publish = random.choice(["none", "rollover"])
signer.dnssec(zone).rrsig_lifetime = 6
signer.dnssec(zone).rrsig_pre_refresh = 1
# options without any effect
signer.dnssec(zone).dnskey_ttl = int(NONSENSE / 10)
signer.dnssec(zone).zone_max_ttl = NONSENSE
signer.dnssec(zone).ksk_lifetime = NONSENSE * 2
signer.dnssec(zone).propagation_delay = int(NONSENSE / 10)

# needed for keymgr
t.generate_conf()

def tickf(when):
    return "+%d" % (STARTUP + when * TICK)

# generate keys, including manual KSK rollover on the beginning
key_ksk1 = signer.key_gen(ZONE, ksk="true", created="+0", publish="+0", ready="+0", active="+0", retire=tickf(4), remove=tickf(5))
key_ksk2 = signer.key_gen(ZONE, ksk="true", created="+0", publish=tickf(3), ready=tickf(4), active=tickf(5), retire="+2h", remove="+3h")
key_zsk1 = knot.key_gen(ZONE, ksk="false", created="+0", publish="+0", active="+0")

# pregenerate keys, exchange KSR, pre-sign it, exchange SKR
KSR = knot.keydir + "/ksr"
SKR = knot.keydir + "/skr"
SKR_BROKEN = SKR + "_broken"
Keymgr.run_check(knot.confile, ZONE, "pregenerate", "+20", "+" + str(FUTURE))
_, out, _ = Keymgr.run_check(knot.confile, ZONE, "generate-ksr", "+0", "+" + str(FUTURE))
writef(KSR, out)
out = signer.sign_ksr(ZONE, KSR)
writef(SKR, out)

cripple_skr(SKR, SKR_BROKEN)
_, _, err = Keymgr.run_check(knot.confile, ZONE, "validate-skr", SKR_BROKEN)
if err.split()[0].casefold() != "error:":
    set_err("keymgr validate-skr")
    detail_log(err)
Keymgr.run_fail(knot.confile, ZONE, "import-skr", SKR_BROKEN)

Keymgr.run_check(knot.confile, ZONE, "import-skr", SKR)

TICK_SAFE = TICK + TICK // 2;

# run it and see if the signing and rollovers work well
if ON_SLAVE:
    master.start()
knot.start()
knot.zone_wait(zone)
check_zone(knot, zone, 2, 1, 1, "init")

zone_update(master, knot, zone, ON_SLAVE)
wait_for_dnskey_count(t, knot, 3, STARTUP + TICK_SAFE * 2)
check_zone(knot, zone, 3, 1 if SIGNER_BIND else 2, 1, "KSK rollover: publish")

zone_update(master, knot, zone, ON_SLAVE)
wait_for_dnskey_count(t, knot, 2, TICK_SAFE * 3)
check_zone(knot, zone, 2, 1, 1, "KSK rollover: finished")

zone_update(master, knot, zone, ON_SLAVE)
wait_for_dnskey_count(t, knot, 3, TICK_SAFE * 2)
check_zone(knot, zone, 3, 1, 1, "ZSK rollover: running")

zone_update(master, knot, zone, ON_SLAVE)
wait_for_dnskey_count(t, knot, 2, TICK_SAFE * 2)
check_zone(knot, zone, 2, 1, 1, "ZSK rollover: done")

# re-generate keys, re-eschange KSR and SKR and re-import it over previous

STARTUP = 1
signer.key_set(ZONE, key_ksk2, retire=tickf(3), remove=tickf(4))
key_ksk3 = signer.key_gen(ZONE, ksk="true", created="+0", publish=tickf(1), ready=tickf(2), active=tickf(3), retire="+4h", remove="+5h")

knot.dnssec(zone).zsk_lifetime = 8 * TICK
knot.gen_confile()

KSR = KSR + "2"
SKR = SKR + "2"
Keymgr.run_check(knot.confile, ZONE, "pregenerate", "+" + str(FUTURE))
_, out, _ = Keymgr.run_check(knot.confile, ZONE, "generate-ksr", "+0", "+" + str(FUTURE))
writef(KSR, out)
out = signer.sign_ksr(ZONE, KSR)
writef(SKR, out)
Keymgr.run_check(knot.confile, ZONE, "import-skr", SKR)

knot.ctl("zone-keys-load")

zone_update(master, knot, zone, ON_SLAVE)
check_zone(knot, zone, 2, 1, 1, "init2")

zone_update(master, knot, zone, ON_SLAVE)
wait_for_dnskey_count(t, knot, 3, STARTUP + TICK_SAFE)
check_zone(knot, zone, 3, 1 if SIGNER_BIND else 2, 1, "KSK rollover2: publish")

zone_update(master, knot, zone, ON_SLAVE)
wait_for_dnskey_count(t, knot, 2, TICK_SAFE * 3)
check_zone(knot, zone, 2, 1, 1, "KSK rollover2: finished")

zone_update(master, knot, zone, ON_SLAVE)
wait_for_dnskey_count(t, knot, 3, TICK_SAFE * 3)
check_zone(knot, zone, 3, 1, 1, "ZSK rollover2: running")

zone_update(master, knot, zone, ON_SLAVE)
wait_for_dnskey_count(t, knot, 2, TICK_SAFE * 2)
check_zone(knot, zone, 2, 1, 1, "ZSK rollover2: done")

# prepare algorithm roll-over: delete pre-generated ZSKs, arrange all the timestamps

if not SIGNER_BIND:
    _, out, _ = Keymgr.run_check(knot.confile, ZONE, "list")
    for line in out.split('\n'):
        if len(line) > 0 and line.split()[-1] == "remove=0": # only one key with this property
            last_zsk = line.split()[0]

    algtick = 5
    now = int(time.time())
    preactive = now + algtick
    publish = preactive + algtick
    postactive = publish + algtick
    remove = postactive + algtick

    key_ksk4 = signer.key_gen(ZONE, ksk="true", algorithm="ECDSAP256SHA256", created="+0",
               pre_active=str(preactive), publish=str(publish), ready=str(publish), active=str(postactive))
    key_zsk2 = knot.key_gen(ZONE, ksk="false", algorithm="ECDSAP256SHA256", created="+0",
               pre_active=str(preactive), publish=str(publish), active=str(postactive))
    signer.key_set(ZONE, key_ksk3, post_active=str(postactive), remove=str(remove))
    knot.key_set(ZONE, last_zsk, post_active=str(postactive), remove=str(remove))

    KSR = KSR + "3"
    SKR = SKR + "3"
    _, out, _ = Keymgr.run_check(knot.confile, ZONE, "generate-ksr", "+0", str(remove + 1))
    writef(KSR, out)
    out = signer.sign_ksr(ZONE, KSR)
    writef(SKR, out)
    Keymgr.run_check(knot.confile, ZONE, "import-skr", SKR)
    knot.ctl("zone-keys-load")

    zone_update(master, knot, zone, ON_SLAVE)
    wait_for_rrsig_count(t, knot, "SOA", 2, algtick + 2)
    check_zone(knot, zone, 2, 1, 2, "alg roll: pre-active")

    zone_update(master, knot, zone, ON_SLAVE)
    wait_for_dnskey_count(t, knot, 4, algtick + 2)
    check_zone(knot, zone, 4, 2, 2, "alg roll: published")

    zone_update(master, knot, zone, ON_SLAVE)
    wait_for_dnskey_count(t, knot, 2, algtick + 2)
    check_zone(knot, zone, 2, 2, 2, "alg roll: post-active")

    zone_update(master, knot, zone, ON_SLAVE)
    wait_for_rrsig_count(t, knot, "SOA", 1, algtick + 2)
    check_zone(knot, zone, 2, 1, 1, "alg roll: finished")

t.end()
