EnterpriseAppProtection/content.js

// content.js

// Encryption utilities
const encryption = {
    async encrypt(data) {
        try {
            const encoder = new TextEncoder();
            const dataBuffer = encoder.encode(JSON.stringify(data));
            const hashBuffer = await crypto.subtle.digest('SHA-256', dataBuffer);
            const hashArray = Array.from(new Uint8Array(hashBuffer));
            return hashArray.map(b => b.toString(16).padStart(2, '0')).join('');
        } catch (error) {
            console.error('Encryption error:', error);
            return null;
        }
    },
    
    async validateData(data) {
        if (!data) return false;
        try {
            const hash = await this.encrypt(data);
            return !!hash;
        } catch (error) {
            console.error('Validation error:', error);
            return false;
        }
    }
};

// Cache management
const cacheManager = {
    cache: new Map(),
    maxSize: 1000,
    maxAge: 3600000, // 1 hour

    set(key, value) {
        if (this.cache.size >= this.maxSize) {
            const oldestKey = this.cache.keys().next().value;
            this.cache.delete(oldestKey);
        }
        this.cache.set(key, {
            value,
            timestamp: Date.now()
        });
    },

    get(key) {
        const item = this.cache.get(key);
        if (!item) return null;
        if (Date.now() - item.timestamp > this.maxAge) {
            this.cache.delete(key);
            return null;
        }
        return item.value;
    },

    clear() {
        this.cache.clear();
    }
};

// Performance monitoring
const performanceMonitor = {
    metrics: {
        linkChecks: 0,
        apiCalls: 0,
        processingTime: 0
    },

    startTimer() {
        return performance.now();
    },

    endTimer(startTime) {
        return performance.now() - startTime;
    },

    logMetric(name, value) {
        this.metrics[name] = (this.metrics[name] || 0) + value;
    },

    getMetrics() {
        return { ...this.metrics };
    },

    reset() {
        this.metrics = {
            linkChecks: 0,
            apiCalls: 0,
            processingTime: 0
        };
    }
};

let domainsDB = {};
let trustedDomains = [];
let blockedDomains = [];
let warningTemplate = "Warning: This link claims to be {app} but goes to an unofficial domain.";
let isDomainsLoaded = false;

// Add rate limiting for API calls
const rateLimiter = {
    lastCall: 0,
    minInterval: 1000, // 1 second between calls
    canMakeCall() {
        const now = Date.now();
        if (now - this.lastCall >= this.minInterval) {
            this.lastCall = now;
            return true;
        }
        return false;
    }
};

// ✅ Get the top-level domain (TLD)
function getTopLevelDomain(hostname) {
    const domainParts = hostname.split(".");
    const knownTLDs = ["co.uk", "com.au", "gov.uk", "edu.au"];

    if (domainParts.length > 2) {
        const lastTwoParts = domainParts.slice(-2).join(".");
        if (knownTLDs.includes(lastTwoParts)) {
            return domainParts.slice(-3).join(".");
        }
    }
    return domainParts.slice(-2).join(".");
}

// ✅ Check if a domain is valid
function isValidDomain(domain, validDomains, trustedDomains) {
    const extractedTLD = getTopLevelDomain(domain);

    if (trustedDomains.includes(domain) || trustedDomains.includes(extractedTLD)) {
        return true; // User has explicitly marked this domain as safe
    }

    return validDomains.some(validDomain => {
        const validTLD = getTopLevelDomain(validDomain);
        return extractedTLD === validTLD;
    });
}

// ✅ Check Google Safe Browsing API for dangerous links
async function checkGoogleSafeBrowsing(url) {
    if (!rateLimiter.canMakeCall()) {
        console.warn("⚠️ Rate limit exceeded for Safe Browsing API calls");
        return false;
    }

    try {
        const response = await chrome.runtime.sendMessage({
            action: "checkSafeBrowsing",
            url: url
        });
        return response.isUnsafe || false;
    } catch (error) {
        console.error("Error checking Safe Browsing:", error);
        return false;
    }
}

// Sanitize warning template to prevent XSS
function sanitizeWarningTemplate(template, appName) {
    const div = document.createElement('div');
    div.textContent = template.replace("{app}", appName);
    return div.textContent;
}

// ✅ Scan page content for suspicious links
async function analyzePageContent() {
    const startTime = performanceMonitor.startTimer();
    
    if (!isDomainsLoaded || !domainsDB || Object.keys(domainsDB).length === 0) {
        console.warn("⚠️ domainsDB is not ready yet. Skipping analysis.");
        return;
    }

    const links = document.querySelectorAll("a:not(.checked-by-extension)");
    let newFlaggedLinks = new Set();
    let processedLinks = 0;

    // Process links in batches of 50
    const batchSize = 50;
    const batches = Array.from(links).reduce((acc, link, i) => {
        const batchIndex = Math.floor(i / batchSize);
        if (!acc[batchIndex]) acc[batchIndex] = [];
        acc[batchIndex].push(link);
        return acc;
    }, []);

    for (const batch of batches) {
        await Promise.all(batch.map(async link => {
            try {
                // Check cache first
                const cachedResult = cacheManager.get(link.href);
                if (cachedResult) {
                    if (cachedResult.isUnsafe) {
                        newFlaggedLinks.add(link.href);
                        addWarningToLink(link, cachedResult.appName);
                    }
                    return;
                }

                // Validate the URL
                let url;
                try {
                    url = new URL(link.href);
                } catch (e) {
                    console.warn("⚠️ Skipping invalid URL:", link.href);
                    return;
                }

                const domain = url.hostname.toLowerCase();
                const linkText = link.innerText.trim();
                processedLinks++;

                // Skip trusted domains
                if (trustedDomains.includes(domain) || trustedDomains.includes(getTopLevelDomain(domain))) {
                    return;
                }

                let matchedApp = null;
                for (const [appName, validDomains] of Object.entries(domainsDB)) {
                    const regex = new RegExp(`\\b${appName}\\b`, "i");
                    if (regex.test(linkText)) {
                        matchedApp = appName;
                    }
                }

                if (matchedApp && domainsDB[matchedApp]) {
                    const validDomains = domainsDB[matchedApp];
                    const isValid = isValidDomain(domain, validDomains, trustedDomains);

                    if (!isValid) {
                        newFlaggedLinks.add(url.href);
                        const isUnsafe = await checkGoogleSafeBrowsing(url.href);
                        
                        // Cache the result
                        cacheManager.set(link.href, {
                            isUnsafe,
                            appName: matchedApp
                        });

                        addWarningToLink(link, matchedApp, isUnsafe);
                    }
                }
            } catch (e) {
                console.error("Unexpected error processing link:", e);
            }
        }));

        // Add a small delay between batches to prevent UI blocking
        await new Promise(resolve => setTimeout(resolve, 50));
    }

    // Update metrics
    performanceMonitor.logMetric('linkChecks', processedLinks);
    performanceMonitor.logMetric('processingTime', performanceMonitor.endTimer(startTime));

    // Store flagged links with encryption
    if (newFlaggedLinks.size > 0) {
        const storageData = {
            flaggedLinks: Array.from(newFlaggedLinks),
            timestamp: Date.now()
        };

        if (await encryption.validateData(storageData)) {
            chrome.storage.local.get(["flaggedLinks", "totalSuspiciousLinks"], async function (result) {
                let existingLinks = new Set(result.flaggedLinks || []);
                let totalCount = existingLinks.size;

                let newLinksToAdd = [...newFlaggedLinks].filter(link => !existingLinks.has(link));

                if (newLinksToAdd.length > 0) {
                    totalCount += newLinksToAdd.length;

                    const updatedData = {
                        totalSuspiciousLinks: totalCount,
                        flaggedLinks: [...existingLinks, ...newLinksToAdd],
                        lastUpdate: Date.now()
                    };

                    if (await encryption.validateData(updatedData)) {
                        chrome.storage.local.set(updatedData, () => {
                            chrome.runtime.sendMessage({
                                action: "updatePopup",
                                flaggedLinks: [...existingLinks, ...newLinksToAdd],
                                totalCount,
                                metrics: performanceMonitor.getMetrics()
                            });
                        });
                    }
                }
            });
        }
    }
}

// Helper function to add warning to link
function addWarningToLink(link, appName, isUnsafe = false) {
    const warning = document.createElement("div");
    warning.classList.add("warning-alert");
    warning.style.cssText = `
        background: #fff3cd;
        color: #856404;
        padding: 10px;
        border: 1px solid #ffeeba;
        border-radius: 4px;
        margin: 5px 0;
        font-size: 14px;
    `;

    warning.textContent = isUnsafe
        ? `⚠️ This link is confirmed dangerous by Google Safe Browsing!`
        : `⚠️ ${sanitizeWarningTemplate(warningTemplate, appName)}`;
    
    warning.setAttribute("title", `This link goes to ${link.hostname} instead of an official ${appName} domain.`);
    link.parentElement.insertBefore(warning, link.nextSibling);
}

// ✅ Load domains and start analysis once ready
function loadDomainsAndAnalyze() {
    chrome.storage.local.get(["domainsDB", "trustedDomains", "blockedDomains", "warningTemplate"], function (result) {
        domainsDB = result.domainsDB || {};
        trustedDomains = result.trustedDomains || [];
        blockedDomains = result.blockedDomains || [];
        warningTemplate = result.warningTemplate || "Warning: This link claims to be {app} but goes to an unofficial domain.";
        console.log("✅ Domains database loaded:", domainsDB);
        isDomainsLoaded = true;
        analyzePageContent();
    });
}

// Debounce function to limit how often a function can be called
function debounce(func, wait) {
    let timeout;
    return function executedFunction(...args) {
        const later = () => {
            clearTimeout(timeout);
            func(...args);
        };
        clearTimeout(timeout);
        timeout = setTimeout(later, wait);
    };
}

// Enhanced error handling wrapper
function withErrorHandling(fn, errorMessage) {
    return async (...args) => {
        try {
            return await fn(...args);
        } catch (error) {
            console.error(`${errorMessage}:`, error);
            // You could add error reporting service here
            return null;
        }
    };
}

// Debounced version of analyzePageContent
const debouncedAnalyzePageContent = debounce(analyzePageContent, 250);

// Enhanced version of analyzePageContent with better error handling
const enhancedAnalyzePageContent = withErrorHandling(debouncedAnalyzePageContent, "Error analyzing page content");

// Update the observer to use the enhanced version
function observePageForLinks() {
    const observer = new MutationObserver(() => enhancedAnalyzePageContent());
    window.pageObserver = observer; // Store for cleanup

    observer.observe(document.body, { childList: true, subtree: true });

    function observeIframes() {
        document.querySelectorAll("iframe").forEach((iframe) => {
            try {
                const iframeDoc = iframe.contentDocument || iframe.contentWindow.document;
                if (iframeDoc) {
                    observer.observe(iframeDoc.body, { childList: true, subtree: true });
                }
            } catch (e) {
                console.warn("Cannot access iframe due to cross-origin restrictions:", iframe.src);
            }
        });
    }

    window.iframeInterval = setInterval(observeIframes, 3000);
}

// ✅ Initialize extension scanning
document.addEventListener("DOMContentLoaded", loadDomainsAndAnalyze);
window.addEventListener("load", loadDomainsAndAnalyze);
observePageForLinks();

// Add cleanup for observers
function cleanup() {
    if (window.pageObserver) {
        window.pageObserver.disconnect();
    }
    if (window.iframeInterval) {
        clearInterval(window.iframeInterval);
    }
}

// Add cleanup on page unload
window.addEventListener("unload", cleanup);