package com.verisignlabs.dnssec.security;

import java.io.IOException;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.SortedMap;
import java.util.TreeMap;
import java.util.logging.Logger;

import org.xbill.DNS.Master;
import org.xbill.DNS.NSEC3PARAMRecord;
import org.xbill.DNS.NSEC3Record;
import org.xbill.DNS.Name;
import org.xbill.DNS.RRSIGRecord;
import org.xbill.DNS.RRset;
import org.xbill.DNS.Record;
import org.xbill.DNS.SOARecord;
import org.xbill.DNS.Type;

public class ZoneData {
    private SortedMap<Name, Set<Integer>> mNodeMap;
    private HashMap<String, RRset> mRRsetMap;
    private SortedMap<Name, MarkRRset> mNSECMap;
    private SortedMap<Name, MarkRRset> mNSEC3Map;
    private Name mZoneName;
    private DNSSECType mDNSSECType;
    private NSEC3PARAMRecord mNSEC3params;
    // Configuration parameters
    private boolean mIgnoreDuplicateRRs;
    private boolean mStripDNSSEC;

    private Logger log = Logger.getLogger(ZoneData.class.toString());

    // The various types of signed zones.
    enum DNSSECType {
        UNSIGNED, NSEC, NSEC3, NSEC3_OPTOUT;
    }

    // The types of nodes (a node consists of all RRs with the same name).
    enum NodeType {
        NORMAL, DELEGATION, GLUE, DNAME, INVALID;
    }

    // if mStripDNSSEC is true, then RRsets of the following types will be skipped
    // on load, or stripped.
    private static final Set<Integer> GENTYPES_SET = Set.of(Type.RRSIG, Type.NSEC, Type.NSEC3, Type.NSEC3PARAM);

    /**
     * This is a subclass of {@link org.xbill.DNS.RRset} that adds a "mark".
     */
    public class MarkRRset extends RRset {

        private static final long serialVersionUID = 1L;
        private boolean mIsMarked = false;

        @Override
        public boolean equals(Object o) {
            return super.equals(o);
        }

        @Override
        public int hashCode() {
            return super.hashCode();
        }

        boolean getMark() {
            return mIsMarked;
        }

        void setMark(boolean value) {
            mIsMarked = value;
        }
    }

    public ZoneData() {
        mNodeMap = new TreeMap<>();
        mRRsetMap = new HashMap<>();
        mNSECMap = new TreeMap<>();
        mNSEC3Map = new TreeMap<>();
    }

    public ZoneData(boolean stripDNSSEC, boolean ignoreDuplicates) {
        super();
        setStripDNSSEC(stripDNSSEC);
        setIgnoreDuplicateRRs(ignoreDuplicates);
    }

    public void setStripDNSSEC(boolean value) {
        this.mStripDNSSEC = value;
    }

    public void setIgnoreDuplicateRRs(boolean value) {
        this.mIgnoreDuplicateRRs = value;
    }

    /**
     * @return The detected DNSSEC "type" (what form of negative proof it uses), or
     *         UNSIGNED if the zone is unsigned or nothing has been processed yet.
     */
    public DNSSECType getDNSSECType() {
        return mDNSSECType;
    }

    /**
     * @return Any detected NSEC3PARAM record in the processed zone, or null if
     *         there wasn't one or no zone has been processed yet.
     */
    public NSEC3PARAMRecord getNSEC3Params() {
        return mNSEC3params;
    }

    public Name zoneName() {
        return mZoneName;
    }

    public SOARecord getSOA() {
        assert mRRsetMap != null;

        RRset r = mRRsetMap.get(key(mZoneName, Type.SOA));
        assert r != null;
     
        return (SOARecord) r.first();
    }
    /**
     * @return The computed "node map" -- this is a mapping of
     *         {@link org.xbill.DNS.Name} objects to simple sets of DNS typecodes.
     *         This will assert if no zone has been processed yet.
     */
    public SortedMap<Name, Set<Integer>> nodeMap() {
        assert mNodeMap != null;
        return mNodeMap;
    }

    /**
     * @return The computed NSEC map.
     */
    public SortedMap<Name, MarkRRset> getNSECMap() {
        assert mNSECMap != null;

        return mNSECMap;
    }

    /**
     * @return the computed NSEC3 map
     */
    public SortedMap<Name, MarkRRset> getNSEC3Map() {
        assert mNSEC3Map != null;

        return mNSEC3Map;
    }

    /**
     * @return an RRset from our internal map with the given name and type, or null if not found.
     */
    public RRset getRRset(Name n, int type) {
        assert mRRsetMap != null;

        return mRRsetMap.get(key(n, type));
    }

    /**
     * Formulate a unique key for an RRset in the RRsetMap. Since an RRset is
     * defined to be the set of Records that all the same ownername and type, we can
     * formulate a valid key by combining them.
     * 
     * @param n    the ownername of the RRset
     * @param type The (integer) type of the RRset
     * @return the key, as a string.
     */
    public static String key(Name n, int type) {
        return n.toString() + ':' + type;
    }

    /**
     * Add an Record (RRSIG or otherwise) to a given RRset
     * 
     * @param rrset The rrset to add to
     * @param rr    The record to add
     * @return true if the record was actually added, false if not.
     */
    private boolean addRRtoRRset(RRset rrset, Record rr) {
        if (mIgnoreDuplicateRRs) {
            rrset.addRR(rr);
            return true;
        }

        if (rr instanceof RRSIGRecord) {
            for (RRSIGRecord sigrec : rrset.sigs()) {
                if (rr.equals(sigrec)) {
                    return false;
                }
            }
        } else {
            for (Record rec : rrset.rrs()) {
                if (rr.equals(rec)) {
                    return false;
                }
            }
        }

        rrset.addRR(rr);
        return true;
    }

    /**
     * Add a record to the various maps.
     * 
     * @param r the record to add
     * @return true if the RR was added, false if it wasn't (because it was a
     *         duplicate)
     */
    private boolean addRR(Record r) {
        assert mNSEC3Map != null;
        assert mRRsetMap != null;

        Name n = r.getName();
        int t = r.getType();
        if (t == Type.RRSIG)
            t = ((RRSIGRecord) r).getTypeCovered();

        // Add NSEC and NSEC3 RRs to their respective maps
        if (t == Type.NSEC) {
            MarkRRset rrset = mNSECMap.computeIfAbsent(n, k -> new MarkRRset());
            return addRRtoRRset(rrset, r);
        }

        if (t == Type.NSEC3) {
            MarkRRset rrset = mNSEC3Map.computeIfAbsent(n, k -> new MarkRRset());
            return addRRtoRRset(rrset, r);
        }

        // Add the name and type to the node map
        Set<Integer> typeset = mNodeMap.computeIfAbsent(n, k -> new HashSet<>());
        typeset.add(r.getType()); // add the original type

        // Add the record to the RRset map
        String k = key(n, t);
        RRset rrset = mRRsetMap.computeIfAbsent(k, k2 -> new RRset());
        return addRRtoRRset(rrset, r);
    }

    /**
     * Given a record, determine the DNSSEC signing type. If the record doesn't
     * determine that, DNSSECType.UNSIGNED is returned
     * 
     * @param r the record to examine
     * @return a DNSSECType of UNKNOWN, unless r is an NSEC or NSEC3 record, in
     *         which case, the type corresponding to the
     */
    private DNSSECType determineDNSSECType(Record r) {
        if (r.getType() == Type.NSEC)
            return DNSSECType.NSEC;
        if (r.getType() == Type.NSEC3) {
            NSEC3Record nsec3 = (NSEC3Record) r;
            if ((nsec3.getFlags() & NSEC3Record.Flags.OPT_OUT) == NSEC3Record.Flags.OPT_OUT) {
                return DNSSECType.NSEC3_OPTOUT;
            }
            return DNSSECType.NSEC3;
        }

        return DNSSECType.UNSIGNED;
    }

    /**
     * Add a resource record to the node map and rrsets map.
     * @param r The record
     * @return the number of errors encountered.
     */
    private int addRecord(Record r) {
        int errors = 0;
        Name n = r.getName();
        int t = r.getType();

        // skip any generated DNSSEC records if we are skipping
        if (mStripDNSSEC && GENTYPES_SET.contains(t)) {
            return errors;
        }

        // Add the record to the various maps.
        boolean res = addRR(r);
        if (!res) {
            log.warning("Record '" + r + "' detected as a duplicate");
            errors++;
        }

        // Learn some things about the zone as we do this pass.
        if (t == Type.SOA)
            mZoneName = n;
        if (t == Type.NSEC3PARAM)
            mNSEC3params = (NSEC3PARAMRecord) r;
        if (mDNSSECType == DNSSECType.UNSIGNED)
            mDNSSECType = determineDNSSECType(r);

        return errors;
    }

    /**
     * Given a zone file, load the node and rrset maps, as well as
     * determine the NSEC3 parameters and signing type.
     *
     * @param master A {@link org.xbill.DNS.Master} object.
     * @return the number of errors encountered.
     * @throws IOException
     */
    public int calculateNodes(String zoneFileName, Name origin) {
        // The zone is unsigned until we get a clue otherwise.
        mDNSSECType = DNSSECType.UNSIGNED;

        int errors = 0;
        try (Master m = zoneFileName.equals("-") ? new Master(System.in) : new Master(zoneFileName, origin)) {
            Record r = null;
            while ((r = m.nextRecord()) != null) {
                errors += addRecord(r);
            }
        } catch (IOException e) {
            e.printStackTrace();
            errors++;
        }

        return errors;
    }

    /**
     * Given an unsorted list of records, load the node and rrset maps, as well as
     * determine the NSEC3 parameters and signing type.
     * 
     * @param records A list of {@link org.xbill.DNS.Record} objects, unsorted.
     * @return the number of errors encountered.
     */
    public int calculateNodes(List<Record> records) {
        mNodeMap = new TreeMap<>();
        mRRsetMap = new HashMap<>();

        // The zone is unsigned until we get a clue otherwise.
        mDNSSECType = DNSSECType.UNSIGNED;

        int errors = 0;
        for (Record r : records) {
            errors += addRecord(r);
        }

        return errors;
    }

    /**
     * Given a name, typeset, and name of the last zone cut, determine the node
     * type.
     * 
     * @param n       The owner name for the node (a collection of RRsets with all
     *                the same owner)
     * @param typeset The (simple) set of types present at that node
     * @param lastCut The name of the last "zone cut". This is the last name (if
     *                encountering the name in DNSSEC canonical name order) that was
     *                either a delegation or a DNAME.
     */
    public NodeType determineNodeType(Name n, Set<Integer> typeset, Name lastCut, Name lastDname) {
        // Names not at or below the zone name are "invalid" -- they shouldn't
        // be in the zone at all, normally. In practice, they are the same as
        // glue.
        if (!n.subdomain(mZoneName)) {
            return NodeType.INVALID;
        }

        // All RRs at the zone apex are normal
        if (n.equals(mZoneName))
            return NodeType.NORMAL;

        // If the node is not below the zone itself, we will treat it as glue (it is
        // really junk).
        if (!n.subdomain(mZoneName)) {
            return NodeType.GLUE;
        }
        // If the node is below a zone cut (either a delegation or DNAME), it is
        // glue.
        if (lastCut != null && n.subdomain(lastCut) && !n.equals(lastCut)) {
            return NodeType.GLUE;
        }
        if (lastDname != null && n.subdomain(lastDname) && !n.equals(lastDname)) {
            return NodeType.GLUE;
        }

        // If the node has a NS record it is a delegation.
        if (typeset.contains(Integer.valueOf(Type.NS)))
            return NodeType.DELEGATION;

        // If the node has a DNAME record, is similar to a delegation
        if (typeset.contains(Integer.valueOf(Type.DNAME))) {
            return NodeType.DNAME;
        }

        // everything else is normal
        return NodeType.NORMAL;
    }

}
