process_html.py

# Name: process_html.py
# Author: Mark Rumsey
# Date: 10/24/2023
# Description: Given html file as input, simplifies elements and
# labels tokens from trained model

from bs4 import BeautifulSoup
from urllib.parse import urlparse
import re
import json
from lxml import etree
from urllib.parse import urlparse
import langdetect
import markdownify
import argparse
from collections import defaultdict

from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
from torch import cuda
from langchain.text_splitter import RecursiveCharacterTextSplitter

XPATH_JSON_PATH = 'models/xpath_queries_2023-07-03_1.json' # Replace with path to JSON file
TOKENIZER_PATH = 'models/tokenizer/tokenizer/'
MODEL_PATH = 'models/model/model/'
MAX_LEN = 514

INLINE_TAGS = ['a', 'abbr', 'acronym', 'b', 'bdo', 'big', 'br', 'button', 'cite', 'code', 'dfn', 'em', 'i',
                'img', 'input', 'kbd', 'label', 'map', 'object', 'output', 'q', 'samp', 'script', 'select',
                'small', 'span', 'strong', 'sub', 'sup', 'textarea', 'time', 'tt', 'var'] 

REMOVE_TAGS = ['nav', 'header', 'menu', 'aside', 'footer', 'script', 'style']

REMOVE_DIVS = ['nav', 'header', 'menu', 'aside', 'footer', 'head', 'search', 'menu', 'navbar', 'navigation',
                'sidebar', 'foot', 'footer', 'bottom', 'pagination']

KEYWORDS = ['CVE', 'CWE', 'date', 'remote', 'local', 'risk', 'host information', 'proof of concept', 'poc',
            'installation', 'steps', 'reproduce', 'platform', 'code', 'impact', 'product', 'vulnerability',
            'vulnerable', 'vulnerabilities', 'component', 'version']


def parse_html(html):
    """
    Parse an HTML code provided as input
    @return: a parsed HTML (i.e., the DOM object)
    """
    return BeautifulSoup(html, 'html.parser')

def simplify_tree(dom):
    """
    Simplify the tree, by keeping only the block-level elements,
    while removing the inline elements (but keeping its innerText).
    @return a simplified DOM object
    """

    # Create a copy of the DOM object to avoid modifying the original
    simplified_dom = BeautifulSoup(str(dom), 'html.parser') 

    for tag in simplified_dom.find_all(['script', 'style']):
        tag.decompose()

    # Remove unnecessary tags and contents
    for tag in simplified_dom.find_all(REMOVE_TAGS): 
        if not any(keyword.lower() in tag.text.lower() for keyword in KEYWORDS):
            tag.decompose() 

    for element in simplified_dom.find_all():
        if element and element.attrs:
            element_id = element.get('id')
            element_class = element.get('class')

            if element_id or element_class:

                # Check if div_id contains the specific word
                if element_id and any(re.search(r'.*{}\.*'.format(re.escape(word)), element_id) for word in REMOVE_DIVS):
                    if not any(keyword.lower() in element.text.lower() for keyword in KEYWORDS):
                        element.decompose()

                # Check if any div_class contains the specific word
                if element_class and any(any(re.search(r'.*{}\.*'.format(re.escape(word)), cls) for word in REMOVE_DIVS) for cls in element_class):
                    if not any(keyword.lower() in element.text.lower() for keyword in KEYWORDS):
                        element.decompose()

    # Remove white space
    simplified_dom = BeautifulSoup(str(simplified_dom).strip(), 'html.parser') 

    return simplified_dom

def load_xpath_queries_from_json(json_file):
    # dictionary
    with open(json_file, 'r') as file:
        xpath_queries = json.load(file)
    return xpath_queries

def remove_nodes_from_json(initial_dom, domain, xpath_queries):

    parsed_dom = etree.HTML(initial_dom.encode())

    if domain in xpath_queries:
        xpath_query_list = xpath_queries[domain]
        # first pass
        flagged_nodes = set()
        for xpath_query in xpath_query_list:
            nodes = parsed_dom.xpath(xpath_query)
            flagged_nodes.update(nodes)
            
        # second pass
        for node in flagged_nodes:
            if node is not None:
                if node.text is not None and node.text.strip() and any(keyword.lower() in node.text.lower() for keyword in KEYWORDS):
                    continue  # Skip if node contains keywords in its text
                parent = node.getparent()
                if parent is not None: # check it is not the root
                    parent.remove(node)

    # Convert the modified etree.Element back to BeautifulSoup
    modified_dom = BeautifulSoup(etree.tostring(parsed_dom, encoding='unicode'), 'html.parser')

    return modified_dom

def remove_html_tags(dom):

    # converts remaining html to markdown
    md = markdownify.markdownify(str(dom))

    return md

def process_html(html_text, url, error_file):
    domain = urlparse(url).netloc
    json_file = XPATH_JSON_PATH
    xpath_queries = load_xpath_queries_from_json(json_file)

    print('Preprocessing...')

    html_dom = parse_html(html_text)

    if domain not in xpath_queries:
        # Apply ONLY filtering and cleaning rules
        # Calling function on 
        filtered_html = simplify_tree(html_dom)
        final_output = remove_html_tags(filtered_html)

    else:
        # Remove identical nodes
        simplified_dom = remove_nodes_from_json(html_dom, domain, xpath_queries)
        # Apply filtering and cleaning rules
        filtered_html = simplify_tree(simplified_dom)
        final_output = remove_html_tags(filtered_html)

    lang = langdetect.detect(str(final_output))

    # skip if no useful data left
    if re.match('<html[^>]*>\s*</html>', str(final_output)) or str(final_output) == '':
        with open(error_file, 'a') as file:
                file.write(f"No remaining html for {url}\n\n")
    # skip if not english
    elif lang != 'en':
        with open(error_file, 'a') as file:
                file.write(f"{url} in {lang}, not English\n\n")
    else:
        return final_output

def run_model(page_data):

    print('Loading model...')

    # load tokenizer and model
    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH)
    model = AutoModelForTokenClassification.from_pretrained(MODEL_PATH)

    # pass page data through model
    pipe = pipeline(task="token-classification", model=model.to("cpu"), tokenizer=tokenizer, aggregation_strategy="simple")

    text_splitter = RecursiveCharacterTextSplitter.from_huggingface_tokenizer(tokenizer=tokenizer, chunk_size=MAX_LEN, chunk_overlap=0, separators=['\n\n', '\n', '.', ' ', ''])
    texts = text_splitter.split_text(page_data)
    entities = defaultdict(list)

    for text in texts:

        predictions = pipe(text)

        # save dictionary of entity:[words]
        for prediction in predictions:
            entities[prediction['entity_group']].append(prediction['word'])

    return entities

def main():
    
    error_file = 'process_html_errors.txt'

    parser = argparse.ArgumentParser()
    parser.add_argument('html', type=argparse.FileType('r'), help='HTML file to pass through model')
    parser.add_argument('url', help='URL source of HTML file')
    parser.add_argument('outfile', type=argparse.FileType('w'), help='JSON output file path')

    args = parser.parse_args()

    page_data = process_html(args.html, args.url, error_file)
    entities = run_model(page_data)

    json.dump(entities, args.outfile)

def run_from_script(html_content, url):
    """Import friendly function."""
    error_file = 'process_html_errors.txt'

    page_data = process_html(html_content, url, error_file)
    if page_data:
        entities = run_model(page_data)

        return entities



if __name__ == '__main__':
    main()