ESH-Media/src/main/search.ts

import * as fs from 'fs';
import * as path from 'path';
import axios from 'axios';
import * as cheerio from 'cheerio';
import { SiteConfig, SearchResult, SearchResultWithSource } from '../shared/types';
import { ProxyManager } from './proxy';
import { app } from 'electron';

/**
 * Search context object that will be passed to custom search scripts
 */
interface SearchContext {
  query: string;
  siteUrl: string;
  useProxy: boolean;
  axios: typeof axios;
  cheerio: typeof cheerio;
  proxyConfig?: {
    host: string;
    port: number;
  };
}

/**
 * Search all enabled sites using their custom scripts
 */
export async function searchAllSites(
  query: string,
  sites: SiteConfig[],
  proxyManager: ProxyManager
): Promise<SearchResultWithSource[]> {
  const enabledSites = sites.filter((site) => site.enabled);

  const searchPromises = enabledSites.map((site) =>
    searchSite(query, site, proxyManager)
      .then((results) => {
        // Add source information to each result
        return results.map((result) => ({
          ...result,
          source: site.id,
          sourceName: site.name,
        }));
      })
      .catch((error) => {
        console.error(`Error searching ${site.name}:`, error.message);
        return [];
      })
  );

  const results = await Promise.all(searchPromises);

  // Flatten and return all results
  return results.flat();
}

/**
 * Search a single site using its custom script
 */
export async function searchSite(
  query: string,
  site: SiteConfig,
  proxyManager: ProxyManager
): Promise<SearchResult[]> {
  try {
    // Load the custom search script
    const scriptPath = getScriptPath(site.searchScript);

    if (!fs.existsSync(scriptPath)) {
      console.error(`Search script not found: ${scriptPath}`);
      return [];
    }

    const scriptContent = fs.readFileSync(scriptPath, 'utf-8');

    // Prepare search context
    const context: SearchContext = {
      query,
      siteUrl: site.url,
      useProxy: site.useProxy,
      axios,
      cheerio,
    };

    // Add proxy config if needed
    if (site.useProxy) {
      const proxyConfig = proxyManager.getProxyConfig();
      if (proxyConfig) {
        context.proxyConfig = {
          host: proxyConfig.host,
          port: proxyConfig.port,
        };
      }
    }

    // Execute the custom script in a safe context
    const results = await executeSearchScript(scriptContent, context);

    // Validate results
    if (!Array.isArray(results)) {
      console.error(`Invalid results from ${site.name}: expected array`);
      return [];
    }

    // Validate each result has required fields
    const validResults = results.filter((result) => {
      if (!result.name || !result.url) {
        console.warn(`Invalid result from ${site.name}:`, result);
        return false;
      }
      return true;
    });

    console.log(`Found ${validResults.length} results from ${site.name}`);

    return validResults;
  } catch (error: any) {
    console.error(`Error executing search script for ${site.name}:`, error.message);
    return [];
  }
}

/**
 * Execute custom search script in a controlled environment
 */
async function executeSearchScript(
  scriptContent: string,
  context: SearchContext
): Promise<SearchResult[]> {
  // Create a safe execution context with required libraries
  const AsyncFunction = Object.getPrototypeOf(async function () {}).constructor;

  // Create the search function from script content
  // The script should define a function that returns a Promise<SearchResult[]>
  const searchFunction = new AsyncFunction(
    'query',
    'siteUrl',
    'useProxy',
    'axios',
    'cheerio',
    'proxyConfig',
    `
    ${scriptContent}

    // Call the search function defined in the script
    if (typeof search === 'function') {
      return await search(query, siteUrl, useProxy, axios, cheerio, proxyConfig);
    } else {
      throw new Error('Search function not defined in script');
    }
    `
  );

  // Execute the function with timeout
  const timeoutPromise = new Promise<SearchResult[]>((_, reject) =>
    setTimeout(() => reject(new Error('Search timeout')), 30000)
  );

  const searchPromise = searchFunction(
    context.query,
    context.siteUrl,
    context.useProxy,
    context.axios,
    context.cheerio,
    context.proxyConfig
  );

  return Promise.race([searchPromise, timeoutPromise]);
}

/**
 * Get absolute path to search script
 */
function getScriptPath(scriptName: string): string {
  const userDataPath = app.getPath('userData');
  const scriptsDir = path.join(userDataPath, 'search-scripts');

  // Check in user data directory first
  let scriptPath = path.join(scriptsDir, scriptName);
  if (fs.existsSync(scriptPath)) {
    return scriptPath;
  }

  // Fallback to default scripts in app directory
  scriptPath = path.join(__dirname, '../../search-scripts', scriptName);
  if (fs.existsSync(scriptPath)) {
    return scriptPath;
  }

  return scriptPath; // Return anyway, will be checked by caller
}

/**
 * List all available search scripts
 */
export function listAvailableScripts(): string[] {
  const scripts: string[] = [];

  // Check default scripts
  const defaultScriptsDir = path.join(__dirname, '../../search-scripts');
  if (fs.existsSync(defaultScriptsDir)) {
    const files = fs.readdirSync(defaultScriptsDir);
    scripts.push(...files.filter((f) => f.endsWith('.js')));
  }

  // Check user scripts
  const userDataPath = app.getPath('userData');
  const userScriptsDir = path.join(userDataPath, 'search-scripts');
  if (fs.existsSync(userScriptsDir)) {
    const files = fs.readdirSync(userScriptsDir);
    scripts.push(...files.filter((f) => f.endsWith('.js') && !scripts.includes(f)));
  }

  return scripts;
}