jsdet_browser/

dom.rs

/// Arena-allocated DOM tree for the browser sandbox.
///
/// Not a full browser DOM — this is a security analysis DOM. It supports:
/// - Element creation, attribute manipulation, child append/remove
/// - querySelector/querySelectorAll (CSS selector subset)
/// - innerHTML/outerHTML get/set
/// - Form and input element enumeration
/// - Visibility tracking (display, visibility, opacity)
///
/// All mutations are observable via the observation stream.
///
use std::collections::HashMap;

/// Opaque node identifier. Indexes into the arena.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
pub struct NodeId(pub u32);

impl NodeId {
    pub const ROOT: Self = Self(0);
}

/// A node in the DOM tree.
#[derive(Debug, Clone)]
pub struct Node {
    pub id: NodeId,
    pub kind: NodeKind,
    pub parent: Option<NodeId>,
    pub children: Vec<NodeId>,
    pub attributes: HashMap<String, String>,
    /// Inline style properties.
    pub style: HashMap<String, String>,
    /// CSS classes.
    pub classes: Vec<String>,
}

/// What kind of node this is.
#[derive(Debug, Clone)]
pub enum NodeKind {
    Document,
    Element { tag: String },
    Text { content: String },
    Comment { content: String },
}

impl Node {
    pub fn tag(&self) -> &str {
        match &self.kind {
            NodeKind::Element { tag } => tag,
            NodeKind::Document => "#document",
            NodeKind::Text { .. } => "#text",
            NodeKind::Comment { .. } => "#comment",
        }
    }

    pub fn is_element(&self) -> bool {
        matches!(self.kind, NodeKind::Element { .. })
    }
}

/// The DOM arena. All nodes are owned here; NodeId indexes into the vec.
#[derive(Debug, Clone)]
pub struct Dom {
    nodes: Vec<Node>,
    /// Quick lookup: tag name → node IDs.
    tag_index: HashMap<String, Vec<NodeId>>,
    /// Quick lookup: id attribute → node ID.
    id_index: HashMap<String, NodeId>,
}

impl Default for Dom {
    fn default() -> Self {
        Self::empty()
    }
}

impl Dom {
    /// Create an empty DOM with just the document node.
    pub fn empty() -> Self {
        let root = Node {
            id: NodeId::ROOT,
            kind: NodeKind::Document,
            parent: None,
            children: Vec::new(),
            attributes: HashMap::new(),
            style: HashMap::new(),
            classes: Vec::new(),
        };
        Self {
            nodes: vec![root],
            tag_index: HashMap::new(),
            id_index: HashMap::new(),
        }
    }

    /// Parse HTML into a DOM tree.
    /// Uses a simple tag-soup parser — not spec-compliant, but handles
    /// the HTML that phishing kits actually produce.
    pub fn parse(html: &str) -> Self {
        let mut dom = Self::empty();

        // Create html, head, body structure.
        let html_id = dom.create_element("html");
        dom.append_child(NodeId::ROOT, html_id);
        let head_id = dom.create_element("head");
        dom.append_child(html_id, head_id);
        let body_id = dom.create_element("body");
        dom.append_child(html_id, body_id);

        // Simple tag-soup parsing into body.
        let mut stack = vec![body_id];
        let mut pos = 0;
        let bytes = html.as_bytes();

        while pos < bytes.len() {
            if bytes[pos] == b'<' {
                // Collect any text before this tag.
                let _tag_start = pos;
                pos += 1;

                if pos < bytes.len() && bytes[pos] == b'/' {
                    // Closing tag — pop stack.
                    while pos < bytes.len() && bytes[pos] != b'>' {
                        pos += 1;
                    }
                    if pos < bytes.len() {
                        pos += 1;
                    }
                    if stack.len() > 1 {
                        stack.pop();
                    }
                } else if pos < bytes.len() && bytes[pos] == b'!' {
                    // Comment or DOCTYPE — skip.
                    while pos < bytes.len() && bytes[pos] != b'>' {
                        pos += 1;
                    }
                    if pos < bytes.len() {
                        pos += 1;
                    }
                } else {
                    // Opening tag.
                    let mut tag_end = pos;
                    while tag_end < bytes.len()
                        && bytes[tag_end] != b'>'
                        && bytes[tag_end] != b' '
                        && bytes[tag_end] != b'/'
                    {
                        tag_end += 1;
                    }
                    let tag_name =
                        String::from_utf8_lossy(&bytes[pos..tag_end]).to_ascii_lowercase();

                    if tag_name.is_empty() {
                        pos = tag_end;
                        continue;
                    }

                    // Parse attributes.
                    let mut attrs = HashMap::new();
                    let mut attr_pos = tag_end;
                    while attr_pos < bytes.len() && bytes[attr_pos] != b'>' {
                        // Skip whitespace.
                        while attr_pos < bytes.len()
                            && (bytes[attr_pos] == b' ' || bytes[attr_pos] == b'/')
                        {
                            attr_pos += 1;
                        }
                        if attr_pos >= bytes.len() || bytes[attr_pos] == b'>' {
                            break;
                        }

                        // Attribute name.
                        let name_start = attr_pos;
                        while attr_pos < bytes.len()
                            && bytes[attr_pos] != b'='
                            && bytes[attr_pos] != b' '
                            && bytes[attr_pos] != b'>'
                            && bytes[attr_pos] != b'/'
                        {
                            attr_pos += 1;
                        }
                        let name = String::from_utf8_lossy(&bytes[name_start..attr_pos])
                            .to_ascii_lowercase();

                        if attr_pos < bytes.len() && bytes[attr_pos] == b'=' {
                            attr_pos += 1;
                            // Attribute value.
                            let value = if attr_pos < bytes.len()
                                && (bytes[attr_pos] == b'"' || bytes[attr_pos] == b'\'')
                            {
                                let quote = bytes[attr_pos];
                                attr_pos += 1;
                                let val_start = attr_pos;
                                while attr_pos < bytes.len() && bytes[attr_pos] != quote {
                                    attr_pos += 1;
                                }
                                let v = String::from_utf8_lossy(&bytes[val_start..attr_pos])
                                    .into_owned();
                                if attr_pos < bytes.len() {
                                    attr_pos += 1;
                                }
                                v
                            } else {
                                let val_start = attr_pos;
                                while attr_pos < bytes.len()
                                    && bytes[attr_pos] != b' '
                                    && bytes[attr_pos] != b'>'
                                {
                                    attr_pos += 1;
                                }
                                String::from_utf8_lossy(&bytes[val_start..attr_pos]).into_owned()
                            };
                            if !name.is_empty() {
                                attrs.insert(name, value);
                            }
                        } else if !name.is_empty() {
                            attrs.insert(name, String::new());
                        }
                    }

                    // Skip to after '>'.
                    while attr_pos < bytes.len() && bytes[attr_pos] != b'>' {
                        attr_pos += 1;
                    }
                    if attr_pos < bytes.len() {
                        attr_pos += 1;
                    }
                    pos = attr_pos;

                    // Create element node.
                    let node_id = dom.create_element(&tag_name);
                    for (k, v) in &attrs {
                        dom.set_attribute(node_id, k, v);
                    }

                    let parent = *stack.last().unwrap_or(&body_id);
                    dom.append_child(parent, node_id);

                    // Self-closing tags don't push onto the stack.
                    let self_closing = matches!(
                        tag_name.as_str(),
                        "br" | "hr"
                            | "img"
                            | "input"
                            | "meta"
                            | "link"
                            | "area"
                            | "base"
                            | "col"
                            | "embed"
                            | "source"
                            | "track"
                            | "wbr"
                    );
                    if !self_closing {
                        stack.push(node_id);
                    }

                    // Handle <script> content — collect until </script>.
                    if tag_name == "script" {
                        // External scripts: find and skip the closing </script> tag.
                        // They have no inline content to collect.
                        if attrs.contains_key("src") {
                            let close_tag = b"</script>";
                            while pos + close_tag.len() <= bytes.len() {
                                if bytes[pos..pos + close_tag.len()].eq_ignore_ascii_case(close_tag)
                                {
                                    pos += close_tag.len();
                                    while pos < bytes.len() && bytes[pos] != b'>' {
                                        pos += 1;
                                    }
                                    if pos < bytes.len() {
                                        pos += 1;
                                    }
                                    break;
                                }
                                pos += 1;
                            }
                            if stack.last() == Some(&node_id) {
                                stack.pop();
                            }
                        } else {
                            let content_start = pos;
                            let close_tag = b"</script>";
                            while pos + close_tag.len() <= bytes.len() {
                                if bytes[pos..pos + close_tag.len()].eq_ignore_ascii_case(close_tag)
                                {
                                    break;
                                }
                                pos += 1;
                            }
                            let script_content =
                                String::from_utf8_lossy(&bytes[content_start..pos]).into_owned();
                            if !script_content.trim().is_empty() {
                                let text_id = dom.create_text(&script_content);
                                dom.append_child(node_id, text_id);
                            }
                            // Skip past </script>.
                            if pos + close_tag.len() <= bytes.len() {
                                pos += close_tag.len();
                                while pos < bytes.len() && bytes[pos] != b'>' {
                                    pos += 1;
                                }
                                if pos < bytes.len() {
                                    pos += 1;
                                }
                            }
                            if stack.last() == Some(&node_id) {
                                stack.pop();
                            }
                        } // end else (inline script)
                    } // end if script

                    // Handle <style> similarly.
                    if tag_name == "style" {
                        let content_start = pos;
                        let close_tag = b"</style>";
                        while pos + close_tag.len() <= bytes.len() {
                            if bytes[pos..pos + close_tag.len()].eq_ignore_ascii_case(close_tag) {
                                break;
                            }
                            pos += 1;
                        }
                        let style_content =
                            String::from_utf8_lossy(&bytes[content_start..pos]).into_owned();
                        if !style_content.trim().is_empty() {
                            let text_id = dom.create_text(&style_content);
                            dom.append_child(node_id, text_id);
                        }
                        if pos + close_tag.len() <= bytes.len() {
                            pos += close_tag.len();
                            while pos < bytes.len() && bytes[pos] != b'>' {
                                pos += 1;
                            }
                            if pos < bytes.len() {
                                pos += 1;
                            }
                        }
                        if stack.last() == Some(&node_id) {
                            stack.pop();
                        }
                    }
                }
            } else {
                // Text content.
                let text_start = pos;
                while pos < bytes.len() && bytes[pos] != b'<' {
                    pos += 1;
                }
                let text = String::from_utf8_lossy(&bytes[text_start..pos]).into_owned();
                if !text.trim().is_empty() {
                    let text_id = dom.create_text(&text);
                    let parent = *stack.last().unwrap_or(&body_id);
                    dom.append_child(parent, text_id);
                }
            }
        }

        dom
    }

    /// Create a new element node. Returns its ID.
    pub fn create_element(&mut self, tag: &str) -> NodeId {
        let id = NodeId(self.nodes.len() as u32);
        let tag_lower = tag.to_ascii_lowercase();
        self.nodes.push(Node {
            id,
            kind: NodeKind::Element {
                tag: tag_lower.clone(),
            },
            parent: None,
            children: Vec::new(),
            attributes: HashMap::new(),
            style: HashMap::new(),
            classes: Vec::new(),
        });
        self.tag_index.entry(tag_lower).or_default().push(id);
        id
    }

    /// Create a text node.
    pub fn create_text(&mut self, content: &str) -> NodeId {
        let id = NodeId(self.nodes.len() as u32);
        self.nodes.push(Node {
            id,
            kind: NodeKind::Text {
                content: content.to_string(),
            },
            parent: None,
            children: Vec::new(),
            attributes: HashMap::new(),
            style: HashMap::new(),
            classes: Vec::new(),
        });
        id
    }

    /// Append a child node to a parent.
    pub fn append_child(&mut self, parent: NodeId, child: NodeId) {
        if let Some(old_parent) = self.nodes[child.0 as usize].parent {
            self.nodes[old_parent.0 as usize]
                .children
                .retain(|c| *c != child);
        }
        self.nodes[child.0 as usize].parent = Some(parent);
        self.nodes[parent.0 as usize].children.push(child);
    }

    /// Set an attribute on a node.
    pub fn set_attribute(&mut self, node: NodeId, name: &str, value: &str) {
        let name_lower = name.to_ascii_lowercase();
        self.nodes[node.0 as usize]
            .attributes
            .insert(name_lower.clone(), value.to_string());

        // Maintain id index.
        if name_lower == "id" {
            self.id_index.insert(value.to_string(), node);
        }

        // Maintain class list.
        if name_lower == "class" {
            self.nodes[node.0 as usize].classes =
                value.split_whitespace().map(|s| s.to_string()).collect();
        }
    }

    /// Get an attribute value.
    pub fn get_attribute(&self, node: NodeId, name: &str) -> Option<&str> {
        self.nodes
            .get(node.0 as usize)?
            .attributes
            .get(&name.to_ascii_lowercase())
            .map(|s| s.as_str())
    }

    /// Remove an attribute.
    pub fn remove_attribute(&mut self, node: NodeId, name: &str) {
        let name_lower = name.to_ascii_lowercase();
        if let Some(node) = self.nodes.get_mut(node.0 as usize) {
            if name_lower == "id"
                && let Some(old_id) = node.attributes.get("id")
            {
                self.id_index.remove(old_id);
            }
            node.attributes.remove(&name_lower);
        }
    }

    /// Get a node by ID.
    pub fn get(&self, id: NodeId) -> Option<&Node> {
        self.nodes.get(id.0 as usize)
    }

    /// Get a mutable node by ID.
    pub fn get_mut(&mut self, id: NodeId) -> Option<&mut Node> {
        self.nodes.get_mut(id.0 as usize)
    }

    /// Find element by id attribute.
    pub fn get_element_by_id(&self, id: &str) -> Option<NodeId> {
        self.id_index.get(id).copied()
    }

    /// Find all elements with the given tag name.
    pub fn get_elements_by_tag(&self, tag: &str) -> Vec<NodeId> {
        self.tag_index
            .get(&tag.to_ascii_lowercase())
            .cloned()
            .unwrap_or_default()
    }

    /// Simple CSS selector query. Supports:
    /// - Tag name: `div`
    /// - ID: `#myid`
    /// - Class: `.myclass`
    /// - Attribute: `[name]`, `[name=value]`
    /// - Combinations: `div.myclass`, `input[type=password]`
    pub fn query_selector(&self, selector: &str) -> Option<NodeId> {
        self.query_selector_all(selector).into_iter().next()
    }

    pub fn query_selector_all(&self, selector: &str) -> Vec<NodeId> {
        let mut results = Vec::new();
        for node in &self.nodes {
            if node.is_element() && self.matches_selector(node, selector) {
                results.push(node.id);
            }
        }
        results
    }

    fn matches_selector(&self, node: &Node, selector: &str) -> bool {
        let selector = selector.trim();

        if let Some(id) = selector.strip_prefix('#') {
            return node.attributes.get("id").is_some_and(|v| v == id);
        }

        if let Some(class) = selector.strip_prefix('.') {
            return node.classes.contains(&class.to_string());
        }

        if selector.starts_with('[')
            && let Some(end) = selector.find(']')
        {
            let attr_spec = &selector[1..end];
            if let Some((name, value)) = attr_spec.split_once('=') {
                let value = value.trim_matches('"').trim_matches('\'');
                return node
                    .attributes
                    .get(&name.to_ascii_lowercase())
                    .is_some_and(|v| v == value);
            } else {
                return node
                    .attributes
                    .contains_key(&attr_spec.to_ascii_lowercase());
            }
        }

        // Tag name, possibly with class: `div.myclass`
        if let Some((tag, rest)) = selector.split_once('.') {
            return node.tag() == tag.to_ascii_lowercase()
                && node.classes.contains(&rest.to_string());
        }
        if let Some((tag, rest)) = selector.split_once('[') {
            let tag_match = tag.is_empty() || node.tag() == tag.to_ascii_lowercase();
            if !tag_match {
                return false;
            }
            let attr_selector = format!("[{rest}");
            return self.matches_selector(node, &attr_selector);
        }

        node.tag() == selector.to_ascii_lowercase()
    }

    /// Get all inline scripts from the DOM.
    pub fn extract_inline_scripts(&self) -> Vec<String> {
        let mut scripts = Vec::new();
        for script_id in self.get_elements_by_tag("script") {
            let node = &self.nodes[script_id.0 as usize];
            // Skip external scripts (have src attribute).
            if node.attributes.contains_key("src") {
                continue;
            }
            // Collect text content.
            for child_id in &node.children {
                if let NodeKind::Text { content } = &self.nodes[child_id.0 as usize].kind
                    && !content.trim().is_empty()
                {
                    scripts.push(content.clone());
                }
            }
        }
        scripts
    }

    /// Get external script URLs.
    pub fn extract_script_urls(&self) -> Vec<String> {
        self.get_elements_by_tag("script")
            .iter()
            .filter_map(|id| self.nodes[id.0 as usize].attributes.get("src").cloned())
            .collect()
    }

    /// Find all form elements with their action URLs.
    pub fn extract_form_actions(&self) -> Vec<String> {
        self.get_elements_by_tag("form")
            .iter()
            .filter_map(|id| self.nodes[id.0 as usize].attributes.get("action").cloned())
            .collect()
    }

    /// Find all password input fields.
    pub fn extract_password_inputs(&self) -> Vec<NodeId> {
        self.get_elements_by_tag("input")
            .into_iter()
            .filter(|id| {
                self.nodes[id.0 as usize]
                    .attributes
                    .get("type")
                    .is_some_and(|t| t.eq_ignore_ascii_case("password"))
            })
            .collect()
    }

    /// Total node count.
    pub fn node_count(&self) -> usize {
        self.nodes.len()
    }

    /// Find the body element.
    pub fn body(&self) -> Option<NodeId> {
        self.get_elements_by_tag("body").into_iter().next()
    }

    /// Find the head element.
    pub fn head(&self) -> Option<NodeId> {
        self.get_elements_by_tag("head").into_iter().next()
    }

    /// Replace a node's children with parsed HTML.
    /// This is the bridge-side implementation of `element.innerHTML = "..."`.
    pub fn set_inner_html(&mut self, node: NodeId, html: &str) {
        let idx = node.0 as usize;
        if idx >= self.nodes.len() {
            return; // invalid node ID — ignore silently
        }
        // Remove existing children
        self.nodes[idx].children.clear();

        // Parse the HTML fragment and attach children to the target node
        let fragment = Dom::parse(html);
        // Copy nodes from fragment into our arena, re-rooting under `node`
        if let Some(frag_body) = fragment.body()
            && let Some(frag_node) = fragment.nodes.get(frag_body.0 as usize)
        {
            let child_ids: Vec<NodeId> = frag_node.children.clone();
            for child_id in child_ids {
                let new_id = self.import_node(&fragment, child_id);
                self.nodes[idx].children.push(new_id);
            }
        }
    }

    /// Import a node (and its subtree) from another DOM into this one.
    fn import_node(&mut self, source: &Dom, source_id: NodeId) -> NodeId {
        let src_idx = source_id.0 as usize;
        let Some(source_node) = source.nodes.get(src_idx) else {
            // Invalid source node — create a placeholder
            let id = NodeId(self.nodes.len() as u32);
            self.nodes.push(Node {
                id,
                kind: NodeKind::Text {
                    content: String::new(),
                },
                parent: None,
                children: Vec::new(),
                attributes: HashMap::new(),
                style: HashMap::new(),
                classes: Vec::new(),
            });
            return id;
        };
        let new_id = NodeId(self.nodes.len() as u32);
        let mut new_node = source_node.clone();
        new_node.id = new_id;
        new_node.children = Vec::new();
        self.nodes.push(new_node);

        // Recursively import children
        let child_ids: Vec<NodeId> = source
            .nodes
            .get(src_idx)
            .map(|n| n.children.clone())
            .unwrap_or_default();
        for child_id in child_ids {
            let imported = self.import_node(source, child_id);
            self.nodes[new_id.0 as usize].children.push(imported);
        }

        new_id
    }

    /// Get the inner HTML of a node (serialized children).
    pub fn inner_html(&self, node: NodeId) -> String {
        let n = &self.nodes[node.0 as usize];
        let mut html = String::new();
        for child in &n.children {
            self.serialize_node(*child, &mut html);
        }
        html
    }

    fn serialize_node(&self, id: NodeId, out: &mut String) {
        let node = &self.nodes[id.0 as usize];
        match &node.kind {
            NodeKind::Text { content } => out.push_str(content),
            NodeKind::Comment { content } => {
                out.push_str("<!--");
                out.push_str(content);
                out.push_str("-->");
            }
            NodeKind::Element { tag } => {
                out.push('<');
                out.push_str(tag);
                for (k, v) in &node.attributes {
                    out.push(' ');
                    out.push_str(k);
                    out.push_str("=\"");
                    out.push_str(v);
                    out.push('"');
                }
                out.push('>');
                for child in &node.children {
                    self.serialize_node(*child, out);
                }
                out.push_str("</");
                out.push_str(tag);
                out.push('>');
            }
            NodeKind::Document => {
                for child in &node.children {
                    self.serialize_node(*child, out);
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parse_simple_html() {
        let dom = Dom::parse("<html><body><div id='test'>hello</div></body></html>");
        assert!(dom.get_element_by_id("test").is_some());
        assert_eq!(dom.get_elements_by_tag("div").len(), 1);
    }

    #[test]
    fn parse_form_with_password() {
        let dom = Dom::parse(
            r#"<form action="/login"><input type="text" name="user"><input type="password" name="pass"></form>"#,
        );
        assert_eq!(dom.extract_form_actions(), vec!["/login"]);
        assert_eq!(dom.extract_password_inputs().len(), 1);
    }

    #[test]
    fn parse_inline_script() {
        let dom = Dom::parse("<script>alert(1)</script>");
        let inline = dom.extract_inline_scripts();
        assert_eq!(inline.len(), 1, "inline: {inline:?}");
        assert!(inline[0].contains("alert(1)"));
    }

    #[test]
    fn parse_external_script() {
        let dom = Dom::parse(r#"<script src="ext.js"></script>"#);
        let urls = dom.extract_script_urls();
        assert_eq!(urls, vec!["ext.js"], "urls: {urls:?}");
        // External scripts should have no inline content.
        assert!(dom.extract_inline_scripts().is_empty());
    }

    #[test]
    fn parse_mixed_scripts() {
        let dom = Dom::parse(
            r#"<div><script>var x=1;</script></div><div><script>var y=2;</script></div>"#,
        );
        let inline = dom.extract_inline_scripts();
        assert_eq!(inline.len(), 2, "inline: {inline:?}");
    }

    #[test]
    fn query_selector_by_class() {
        let dom = Dom::parse(r#"<div class="a b"><span class="b c">x</span></div>"#);
        assert_eq!(dom.query_selector_all(".b").len(), 2);
    }

    #[test]
    fn query_selector_by_attribute() {
        let dom = Dom::parse(r#"<input type="password"><input type="text">"#);
        let results = dom.query_selector_all("input[type=password]");
        assert_eq!(results.len(), 1);
    }

    #[test]
    fn create_and_append() {
        let mut dom = Dom::empty();
        let body = dom.create_element("body");
        dom.append_child(NodeId::ROOT, body);
        let div = dom.create_element("div");
        dom.set_attribute(div, "id", "new");
        dom.append_child(body, div);
        assert!(dom.get_element_by_id("new").is_some());
    }

    #[test]
    fn inner_html_serialization() {
        let dom = Dom::parse("<div><span>hello</span></div>");
        let divs = dom.get_elements_by_tag("div");
        assert!(!divs.is_empty());
        let html = dom.inner_html(divs[0]);
        assert!(html.contains("<span>"));
        assert!(html.contains("hello"));
    }
}