<?php
class MalwareTool
{
    const ZIP_PASSWORD = 'infected';
    const ADVANCED_EXTRACTION_SCRIPT_PATH = APP . 'files/scripts/generate_file_objects.py';

    /**
     * @param string $originalFilename
     * @param string $content
     * @param string $md5
     * @return string Content of zipped file
     * @throws Exception
     */
    public function encrypt($originalFilename, $content, $md5)
    {
        if (method_exists("ZipArchive", "setEncryptionName")) {
            // When PHP zip extension is installed and supports creating encrypted archives.
            return $this->encryptByExtension($originalFilename, $content, $md5);
        } else {
            return $this->encryptByCommand($originalFilename, $content, $md5);
        }
    }

    /**
     * @param string $originalFilename
     * @param string $content
     * @param string $md5
     * @return string Content of zipped file
     * @throws Exception
     */
    private function encryptByCommand($originalFilename, $content, $md5)
    {
        $tempDir = $this->tempDir();

        $contentsFile = new File($tempDir . DS . $md5, true);
        if (!$contentsFile->write($content)) {
            throw new Exception("Could not write content to file '{$contentsFile->path}'.");
        }
        $contentsFile->close();

        $fileNameFile = new File($tempDir . DS . $md5 . '.filename.txt', true);
        if (!$fileNameFile->write($originalFilename)) {
            throw new Exception("Could not write original file name to file '{$fileNameFile->path}'.");
        }
        $fileNameFile->close();

        $zipFile = new File($tempDir . DS . $md5 . '.zip');

        $exec = [
            'zip',
            '-j', // junk (don't record) directory names
            '-P', // use standard encryption
            self::ZIP_PASSWORD,
            escapeshellarg($zipFile->path),
            escapeshellarg($contentsFile->path),
            escapeshellarg($fileNameFile->path),
        ];

        try {
            $this->execute($exec);
            $zipContent = $zipFile->read();
            if ($zipContent === false) {
                throw new Exception("Could not read content of newly created ZIP file.");
            }

            return $zipContent;

        } catch (Exception $e) {
            throw new Exception("Could not create encrypted ZIP file '{$zipFile->path}'.", 0, $e);

        } finally {
            $fileNameFile->delete();
            $contentsFile->delete();
            $zipFile->delete();
        }
    }

    /**
     * @param string $originalFilename
     * @param string $content
     * @param string $md5
     * @return string Content of zipped file
     * @throws Exception
     */
    private function encryptByExtension($originalFilename, $content, $md5)
    {
        $zipFilePath = $this->tempFileName();

        $zip = new ZipArchive();
        $result = $zip->open($zipFilePath, ZipArchive::CREATE);
        if ($result === true) {
            $zip->setPassword(self::ZIP_PASSWORD);

            $zip->addFromString($md5, $content);
            $zip->setEncryptionName($md5, ZipArchive::EM_AES_128);

            $zip->addFromString("$md5.filename.txt", $originalFilename);
            $zip->setEncryptionName("$md5.filename.txt", ZipArchive::EM_AES_128);

            $zip->close();
        } else {
            throw new Exception("Could not create encrypted ZIP file '$zipFilePath'. Error code: $result");
        }

        $zipFile = new File($zipFilePath);
        $zipContent = $zipFile->read();
        if ($zipContent === false) {
            throw new Exception("Could not read content of newly created ZIP file.");
        }
        $zipFile->delete();

        return $zipContent;
    }

    /**
     * @param string $content
     * @param array $hashTypes
     * @return array
     * @throws InvalidArgumentException
     */
    public function computeHashes($content, array $hashTypes = array())
    {
        $validHashes = array('md5', 'sha1', 'sha256');
        $hashes = [];
        foreach ($hashTypes as $hashType) {
            if (!in_array($hashType, $validHashes)) {
                throw new InvalidArgumentException("Hash type '$hashType' is not valid hash type.");
            }
            $hashes[$hashType] = hash($hashType, $content);
        }
        return $hashes;
    }

    /**
     * @param string $pythonBin
     * @param string $filePath
     * @return array
     * @throws Exception
     */
    public function advancedExtraction($pythonBin, $filePath)
    {
        return $this->executeAndParseJsonOutput([
            $pythonBin,
            self::ADVANCED_EXTRACTION_SCRIPT_PATH,
            '-p',
            escapeshellarg($filePath),
        ]);
    }

    /**
     * @param string $pythonBin
     * @return array
     * @throws Exception
     */
    public function checkAdvancedExtractionStatus($pythonBin)
    {
        return $this->executeAndParseJsonOutput([$pythonBin, self::ADVANCED_EXTRACTION_SCRIPT_PATH, '-c']);
    }

    private function tempFileName()
    {
        $randomName = (new RandomTool())->random_str(false, 12);
        return $this->tempDir() . DS . $randomName;
    }

    /**
     * @return string
     */
    private function tempDir()
    {
        return Configure::read('MISP.tmpdir') ?: sys_get_temp_dir();
    }

    /**
     * @param array $command
     * @return array
     * @throws Exception
     */
    private function executeAndParseJsonOutput(array $command)
    {
        $output = $this->execute($command);

        $json = json_decode($output, true);
        if ($json === null) {
            throw new Exception("Command output is not valid JSON: " . json_last_error_msg());
        }
        return $json;
    }

    /**
     * This method is much more complicated than just `exec`, but it also provide stderr output, so Exceptions
     * can be much more specific.
     *
     * @param array $command
     * @return string
     * @throws Exception
     */
    private function execute(array $command)
    {
        $descriptorspec = [
            1 => ["pipe", "w"], // stdout
            2 => ["pipe", "w"], // stderr
        ];

        $command = implode(' ', $command);
        $process = proc_open($command, $descriptorspec, $pipes);
        if (!$process) {
            throw new Exception("Command '$command' could be started.");
        }

        $stdout = stream_get_contents($pipes[1]);
        if ($stdout === false) {
            throw new Exception("Could not get STDOUT of command.");
        }
        fclose($pipes[1]);

        $stderr = stream_get_contents($pipes[2]);
        fclose($pipes[2]);

        $returnCode = proc_close($process);
        if ($returnCode !== 0) {
            throw new Exception("Command '$command' return error code $returnCode. STDERR: '$stderr', STDOUT: '$stdout'");
        }

        return $stdout;
    }
}