package net.deterlab.abac;

import org.apache.commons.collections15.*;

import edu.uci.ics.jung.graph.*;
import edu.uci.ics.jung.graph.util.*;

import java.util.*;

/**
 * A class for making queries against the graph. It supports direct queries as
 * well as reachability in either direction. See the run method for details.
 * @author <a href="http://abac.deterlab.net">ISI ABAC team</a>
 * @version 1.3
 */
class Query {
    /** Internal graph representation */
    private Graph<Role,Credential> g;
    /** Count of vertices in g */
    private int vertex_count;

    /**
     * Create a query to be run against the credential graph.
     * @param g a Graph represneting the credentials and implicit connections.
     */
    public Query(Graph<Role,Credential> g) {
        this.g = g;
    }

    /**
     * Run a query against the graph, returning a graph of the results. If the
     * results are empty or the query fails, an empty graph is returned. When
     * derived edges are involved, the subgraphs that imply those edges are
     * included.
     * @param attr a String containing the role to look for
     * @param prin a String containing the principal
     * @return a Graph with the proof or partial proof.
     */
    public Graph<Role,Credential> run(String attr, String prin) {
        Role attribute = null, principal = null;;

        if (!attr.isEmpty()) attribute = new Role(attr);
        if (!prin.isEmpty()) principal = new Role(prin);

        Graph<Role,Credential> ret =
            Graphs.<Role,Credential>synchronizedDirectedGraph(
                    new DirectedSparseGraph<Role,Credential>());

	/* empty attribute, non-empty principal: find everywhere the principal
	 * can go */
        if (attr.isEmpty() && !prin.isEmpty()) {
            if (g.containsVertex(principal)) {
                CollectReversePath collect = new CollectReversePath(ret);
                forward_dfs(principal, collect);
            }
        }

        /* otherwise we're going to do some kind of a reverse dfs */
        else if (g.containsVertex(attribute)) { 
            if ( prin == null || prin.isEmpty()) {
                CollectQueryGraph  collect = new CollectQueryGraph(ret);
                reverse_dfs(attribute, collect);
            }
            else {
                if ( g.containsVertex(principal)) {
                    CollectQueryPath  collect = new CollectQueryPath(ret,
			    principal);
                    reverse_dfs(attribute, collect);
                }
            }
        }

        vertex_count = ret.getVertexCount();

        return ret;
    }

    /**
     * Returns true after running a query that returns a non-empty set of
     * vertices.
     * @return a boolean, true after running a query that returns a non-empty
     * set of vertices.
     */
    public boolean successful() {
        return vertex_count > 0;
    }

    /**
     * Returns a collection of principals reachable from a Role when
     * traversing edges in the reverse direction.
     * @param n the Role to start from
     * @return a Set containinf the principals
     */
    public Set<Role> find_principals(Role n) {
        Set<Role> principals = new HashSet<Role>();
    
        GetPrincipals f = new GetPrincipals(principals);
        reverse_dfs(n, f);
        return principals;
    }


    /**
     * Interface to pass to the reverse_dfs class that operates on nodes in the
     * graph as part of the depth first search (DFS).
     */
    interface DfsFcn {
	/**
	 * Called when the node is first encountered on the DFS.
	 * @param x the Role being visited
	 */
        public void node(Role x);
	/**
	 * Called when all the node's children have been visited on the DFS.
	 * @param x the Role being visited
	 */
        public void node_after(Role x);
    }

    /**
     * Class used to pull all the principals out of a reverse_dfs searched
     * graph.  It takes a collection to the constructor and adds all principals
     * it sees to that collection.  In no collection is passed in, it creates a
     * local HashSet of vertices.  The collection is accessible via the p
     * (i.e., principals) member.
     */
    private class GetPrincipals implements DfsFcn {
	/** The principals collected */
        public Collection<Role> p;

	/**
	 * Initialize the principals collection.
	 * @param vc the initial collection.  May be null.
	 */
        public GetPrincipals(Collection<Role> vc) {
            if ( vc != null) { p =vc;}
            else { p = new HashSet<Role>(); }
        }

	/**
	 * Adds principals to the collection, called by the DFS.
	 * @param r the Role visited.
	 */
        public void node(Role r) {
            if ( r.is_principal() ) p.add(r);
        }
	/**
	 * Does nothing.
	 * @param r a Role, a dummy
	 */
        public void node_after(Role r) { }
    }

    /**
     * Collect a subgraph. Add every node and outgoing arc we encounter, and
     * pull in the linking node and intersection subtrees as well. Because it's
     * possible that these subtrees have been encountered more than once, we're
     * careful not to pull them in twice. The created graph is found in the
     * public ng member.
     */
    private class CollectQueryGraph implements DfsFcn {
	/** The subgraph being collected */
        public Graph<Role, Credential> ng;
	/** The linking roles already added to the sub-graph */
        protected HashSet<Role> linking_roles_seen;
	/** The intersection roles already added to the sub-graph */
        protected HashSet<Role> intersection_roles_seen;

	/**
	 * Initialize the graph into which to gather nodes/edges.
	 * @param g A graph to collect
	 */
        public CollectQueryGraph(Graph<Role,Credential> g) {
            ng = g;
            linking_roles_seen = new HashSet<Role>();
            intersection_roles_seen = new HashSet<Role>();
        }


	/**
	 * Do the collection on each node we visit.
	 * @param r a Role, the current node in the DFS.
	 */
        public void node(Role r) {
            if (r.is_linking() && !linking_roles_seen.contains(r)) {
                // Collect this linking role subgraph
                reverse_dfs(new Role(r.principal_part()), this);
                linking_roles_seen.add(r);
            }

            // collect subgraph for each intersection prereq
            if (r.is_intersection() && !intersection_roles_seen.contains(r)) {
                for (Role prereq : r.prereqs())
                    reverse_dfs(prereq, this);
                intersection_roles_seen.add(r);
            }

            // Add this node's children
            for (Credential c : g.getInEdges(r) ) {
                Role tail = g.getSource(c);
                ng.addEdge(c, tail, r); // FIXME tail and r can come from c [?]
            }
        }
	/**
	 * Does nothing.
	 * @param r a Role, a dummy
	 */
        public void node_after(Role r) { }

    }

    /**
     * Collect the path from one node to another. It includes vertices that
     * are traversed to reach the node. It also includes subgraphs that imply
     * the linking/intersection edges traversed.
     */
    private class CollectQueryPath extends CollectQueryGraph {
	/** Vertices on the path to the principal */
        HashSet<Role> onPath;
	/**
	 * Establish the collection graph and the destination.
	 * @param g the graph that collects the result
	 * @param dest the Role to find
	 */
        public CollectQueryPath(Graph<Role,Credential> g, Role dest) {
            super(g);
            onPath = new HashSet<Role>();
            onPath.add(dest);
        }

	/**
	 * Does nothing.
	 * @param r a Role, a dummy
	 */
        public void node(Role r) { }

	/**
	 * Do the collection on each node we visit.
	 * @param r a Role, the current node in the DFS.
	 */
        public void node_after(Role r) {
            for (Credential c : g.getInEdges(r)) {
                Role child = c.tail();

                if (onPath.contains(child)) {
                    Credential edge = g.findEdge(child, r);
                    if (edge == null)
                        throw new RuntimeException("Credential missing from " +
				"parent graph, state is messed up");

                    ng.addEdge(edge, child, r);
                    onPath.add(r);

                    // For linking roles, collect the whole subgraph of the
                    // authorizer
                    if (r.is_linking()) {
                        CollectQueryPath link = new CollectQueryPath(ng, 
				new Role(child.principal_part()));
                        reverse_dfs(new Role(r.principal_part()), link);
                    }

                    // intersection: collect the subgraph of each prereq
                    else if (r.is_intersection()) {
                        CollectQueryPath prereq_finder = 
			    new CollectQueryPath(ng, child);
                        for (Role prereq : r.prereqs())
                            reverse_dfs(prereq, prereq_finder);
                    }
                }
            }
        }
    }

    /**
     * Collect the path that a node can reach. It includes subgraphs that imply
     * the linked/intersection edges traversed.
     */
    private class CollectReversePath extends CollectQueryGraph {
	/**
	 * Initialize the collection graph.
	 * @param g a Graph that collects the results.
	 */
        public CollectReversePath(Graph<Role,Credential> g) {
            super(g);
        }

        /**
	 * add all of the node's parents
	 * @param r a Role, the current node being visited
	 */
        public void node(Role r) {
            for (Credential c : g.getOutEdges(r)) {
                Role head = g.getDest(c);
                Credential cred = g.findEdge(r, head);
                if (cred == null)
                    throw new RuntimeException("Credential missing from " + 
			    "parent graph, state is messed up");

                ng.addEdge(cred, r, head);
            }
        }

	/**
	 * Do the collection on each node we visit.
	 * @param r a Role, the current node in the DFS.
	 */
        public void node_after(Role r) {
            for (Credential c : g.getOutEdges(r)) {
                Role parent = g.getDest(c);

                // if any of the links we follow is from a linking node, copy
                // the subgraph that implies said link
                if (parent.is_linking()) {
                    CollectQueryPath link = new CollectQueryPath(ng, 
			    new Role(r.principal_part()));
                    reverse_dfs(new Role(parent.principal_part()), link);
                }

                // intersection roles do not need to be treated specially
                // because the forward path from the principal will traverse all
                // edges that imply the edge to the intersection node
            }
        }
    }

    /**
     * Interface to the reverse_dfs member that does allocates the visited map
     * so the user doesn't have to.
     * @param r a Role from which to start the search
     * @param f a DfsFcn that customizes the search operation
     */
    private void reverse_dfs(Role r, DfsFcn f) {
        reverse_dfs(r, new HashSet<Role>(), f);
    }

    /**
     * Member function to walk the edges in reverse from a given Role.  Each
     * Role is visited once and has the node() function of the dfs_fcn called
     * on it.  Order is not guaranteed.
     * @param r a Role from which to start the search
     * @param visited a HashSet of Roles that tells which nodes have been
     * visited.
     * @param f a DfsFcn that customizes the search operation
     */
    private void reverse_dfs(Role r, HashSet<Role> visited,
            DfsFcn f) {
        if (visited.contains(r)) return;

        f.node(r);
        visited.add(r);

        for (Credential c : g.getInEdges(r)) 
            reverse_dfs(g.getSource(c), visited, f);
        f.node_after(r);
    }

    /**
     * Interface to the forward_dfs member that allocates the visited map so
     * the user doesn't have to.
     * @param r a Role from which to start the search
     * @param f a DfsFcn that customizes the search operation
     */
    private void forward_dfs(Role r, DfsFcn f) {
        forward_dfs(r, new HashSet<Role>(), f);
    }

    /**
     * Member function to walk the edges from a given Role. Each Role is
     * visited once and has the node() function of the dfs_fcn called on it.
     * Order is not guaranteed
     * @param r a Role from which to start the search
     * @param visited a HashSet of Roles that tells which nodes have been
     * visited.
     * @param f a DfsFcn that customizes the search operation
     */
    private void forward_dfs(Role r, Set<Role> visited, DfsFcn f) {
        if (visited.contains(r)) return;

        f.node(r);
        visited.add(r);

        for (Credential c : g.getOutEdges(r))
            forward_dfs(g.getDest(c), visited, f);
        f.node_after(r);
    }
}