<?php
namespace Relay\Restore;

use Relay\Database\Db;
use Relay\Support\RelayError;

/**
 * Whitelist for SQL read from a backup archive. Only the statement forms Relay's own writer produces are
 * accepted, INSERT values must be plain literals, and every table name is rewritten from the backup's name
 * to this site's import (or final) name. Anything else is rejected before it reaches the database.
 */
final class SqlGate
{
    private const IDENT = '`((?:[^`]|``)+)`';
    private const COLS = '`(?:[^`]|``)+`(?:, ?`(?:[^`]|``)+`)*';
    private const ACTIONS = '(?: ON (?:DELETE|UPDATE) (?:RESTRICT|CASCADE|SET NULL|NO ACTION|SET DEFAULT))*';
    // MEMORY: tables kept in RAM, such as Wordfence Login Security's wp_wfls_role_counts.
    private const ENGINES = ['innodb', 'myisam', 'aria', 'memory'];

    /** @var array<string, string> backup table => table to write rows into */
    private array $import;
    /** @var array<string, string> backup table => final table name on this site */
    private array $final;
    /** @var string[] lower-case collations this server has; empty = keep collations as written */
    private array $collations;
    /** Whether the SQL comes from another site (a pull or push) rather than a backup: refusals say so. */
    private bool $from_site;

    public function __construct(array $import, array $final, array $collations = [], bool $from_site = false)
    {
        $this->import = $import;
        $this->final = $final;
        $this->collations = $collations;
        $this->from_site = $from_site;
    }

    /** @return string[] statements of one database part (Relay's writer ends each with ";\n" and never emits a raw newline in a value). */
    public static function split(string $sql): array
    {
        $out = [];
        foreach (explode(";\n", $sql) as $statement) {
            $statement = trim((string) preg_replace('/\A(?:--[^\n]*\n)+/', '', $statement));
            if ($statement !== '' && strpos($statement, '--') !== 0) {
                $out[] = $statement;
            }
        }
        return $out;
    }

    /** @return array{kind: string, table: ?string, sql: string} */
    public function check(string $sql): array
    {
        if (preg_match("/^SET (?:NAMES [A-Za-z0-9_]+(?: COLLATE [A-Za-z0-9_]+)?|sql_mode = '[A-Z_,]*'|FOREIGN_KEY_CHECKS = [01]|time_zone = '[+-]\\d{2}:\\d{2}')$/", $sql)) {
            return ['kind' => 'set', 'table' => null, 'sql' => $sql];
        }
        if (preg_match('/^DROP TABLE IF EXISTS ' . self::IDENT . '$/', $sql, $m)) {
            $t = $this->known($m[1], $sql);
            return ['kind' => 'drop', 'table' => $t, 'sql' => 'DROP TABLE IF EXISTS ' . Db::ident($this->import[$t])];
        }
        if (strpos($sql, 'CREATE TABLE ') === 0) {
            return $this->create($sql);
        }
        if (strpos($sql, 'INSERT INTO ') === 0) {
            return $this->insert($sql);
        }
        if (strpos($sql, 'ALTER TABLE ') === 0) {
            return $this->constraint($sql);
        }
        throw $this->rejected($sql); // phpcs:ignore WordPress.Security.EscapeOutput.ExceptionNotEscaped -- raw SQL; rejected() shortens and escapes it in the message
    }

    private function create(string $sql): array
    {
        if (!preg_match('/^CREATE TABLE ' . self::IDENT . ' \(/', $sql, $m)) {
            throw $this->rejected($sql); // phpcs:ignore WordPress.Security.EscapeOutput.ExceptionNotEscaped -- raw SQL; rejected() shortens and escapes it in the message
        }
        $t = $this->known($m[1], $sql);
        $body = self::table_body(substr($sql, strlen($m[0])), $this->collations);
        if ($body === null) {
            throw $this->rejected($sql); // phpcs:ignore WordPress.Security.EscapeOutput.ExceptionNotEscaped -- raw SQL; rejected() shortens and escapes it in the message
        }
        return ['kind' => 'create', 'table' => $t, 'sql' => 'CREATE TABLE ' . Db::ident($this->import[$t]) . ' (' . $body];
    }

    /**
     * The part of a CREATE TABLE after its opening parenthesis, checked: no SELECT, LIKE, foreign keys, other
     * storage or engines, nothing after a semicolon. Collations this server lacks are fitted. Null when refused.
     *
     * @param string[] $collations lower-case collations this server has; empty = keep them as written
     */
    public static function table_body(string $rest, array $collations): ?string
    {
        // Keywords are checked with string literals and identifiers blanked out, so comments and names can't trip (or hide) them.
        $bare = (string) preg_replace(["/'(?:[^'\\\\]++|\\\\.)*+'/s", '/`(?:[^`]++|``)*+`/'], ["''", '``'], $rest);
        if (preg_match('/\b(SELECT|LIKE|DATA\s+DIRECTORY|INDEX\s+DIRECTORY|CONNECTION|UNION|REFERENCES|TABLESPACE)\b|;/i', $bare)) {
            return null;
        }
        preg_match_all('/\bENGINE\s*=?\s*([A-Za-z0-9_]+)/i', $bare, $engines);
        foreach ($engines[1] as $engine) {
            if (!in_array(strtolower($engine), self::ENGINES, true)) {
                return null;
            }
        }
        return self::fit_collations($rest, $bare, $collations);
    }

    /**
     * A collation this server lacks (MySQL 8's utf8mb4_0900_*, MariaDB 11's uca1400) becomes the closest one it
     * has. utf8mb4 is never downgraded: a server without it can't hold the data, so the restore stops instead.
     */
    private static function fit_collations(string $rest, string $bare, array $collations): string
    {
        if ($collations === []) {
            return $rest;
        }
        if (stripos($bare, 'utf8mb4') !== false && preg_grep('/^utf8mb4_/', $collations) === []) {
            throw new RelayError(
                'db.charset',
                esc_html__('This database server does not support utf8mb4 (full Unicode, used for emoji and many languages), and SiteRelay will not convert data to a smaller character set.', 'siterelay'),
                esc_html__('Use a server with MySQL 5.5.3+ or MariaDB 5.5+.', 'siterelay')
            );
        }
        return (string) preg_replace_callback("/'(?:[^'\\\\]++|\\\\.)*+'|\\b(COLLATE)([ =]+)([A-Za-z0-9_]+)/i", static function (array $m) use ($collations): string {
            return isset($m[3]) ? $m[1] . $m[2] . self::collation($m[3], $collations) : $m[0];
        }, $rest);
    }

    private static function collation(string $name, array $collations): string
    {
        $lower = strtolower($name);
        if (in_array($lower, $collations, true)) {
            return $name;
        }
        $charset = strtok($lower, '_');
        foreach (['_unicode_520_ci', '_unicode_ci', '_general_ci', '_swedish_ci'] as $suffix) {
            if (in_array($charset . $suffix, $collations, true)) {
                return $charset . $suffix;
            }
        }
        return $name; // nothing close: the server reports it
    }

    private function insert(string $sql): array
    {
        if (!preg_match('/^INSERT INTO ' . self::IDENT . ' \(/', $sql, $m)) {
            throw $this->rejected($sql); // phpcs:ignore WordPress.Security.EscapeOutput.ExceptionNotEscaped -- raw SQL; rejected() shortens and escapes it in the message
        }
        $t = $this->known($m[1], $sql);
        $p = strlen($m[0]);
        while (preg_match('/\G`(?:[^`]|``)+`(,|\) VALUES )/', $sql, $c, 0, $p)) {
            $p += strlen($c[0]);
            if ($c[1] !== ',') {
                break;
            }
        }
        if (substr($sql, $p - 8, 8) !== ' VALUES ') {
            throw $this->rejected($sql); // phpcs:ignore WordPress.Security.EscapeOutput.ExceptionNotEscaped -- raw SQL; rejected() shortens and escapes it in the message
        }
        $values_at = $p;
        $n = strlen($sql);
        while (true) {
            if ($p >= $n || $sql[$p] !== '(') {
                throw $this->rejected($sql); // phpcs:ignore WordPress.Security.EscapeOutput.ExceptionNotEscaped -- raw SQL; rejected() shortens and escapes it in the message
            }
            $p++;
            while (true) {
                $p = self::value($sql, $p);
                if ($p < 0 || $p >= $n) {
                    throw $this->rejected($sql); // phpcs:ignore WordPress.Security.EscapeOutput.ExceptionNotEscaped -- raw SQL; rejected() shortens and escapes it in the message
                }
                if ($sql[$p] === ',') {
                    $p++;
                    continue;
                }
                if ($sql[$p] !== ')') {
                    throw $this->rejected($sql); // phpcs:ignore WordPress.Security.EscapeOutput.ExceptionNotEscaped -- raw SQL; rejected() shortens and escapes it in the message
                }
                $p++;
                break;
            }
            if ($p === $n) {
                break;
            }
            if ($sql[$p] !== ',') {
                throw $this->rejected($sql); // phpcs:ignore WordPress.Security.EscapeOutput.ExceptionNotEscaped -- raw SQL; rejected() shortens and escapes it in the message
            }
            $p++;
        }
        $columns = substr($sql, strlen($m[0]) - 1, $values_at - strlen($m[0]) + 1);
        return ['kind' => 'insert', 'table' => $t, 'sql' => 'INSERT INTO ' . Db::ident($this->import[$t]) . ' ' . $columns . substr($sql, $values_at)];
    }

    /** Position after one literal value starting at $p, or -1. */
    private static function value(string $s, int $p): int
    {
        $c = $s[$p] ?? '';
        if ($c === "'") {
            $n = strlen($s);
            $p++;
            while (true) {
                $p += strcspn($s, "'\\", $p);
                if ($p >= $n) {
                    return -1;
                }
                if ($s[$p] === '\\') {
                    $p += 2;
                    continue;
                }
                return $p + 1;
            }
        }
        if (substr_compare($s, 'NULL', $p, 4) === 0) {
            return $p + 4;
        }
        $num = '-?\d+(?:\.\d+)?(?:[eE][-+]?\d+)?';
        foreach ([$num, "UNHEX\\('[0-9A-Fa-f]*'\\)", "b'[01]*'", "ST_GeomFromText\\('POINT\\({$num} {$num}\\)'\\)"] as $pattern) {
            if (preg_match('/\G' . $pattern . '/', $s, $m, 0, $p)) {
                return $p + strlen($m[0]);
            }
        }
        return -1;
    }

    private function constraint(string $sql): array
    {
        if (!preg_match('/^ALTER TABLE ' . self::IDENT . ' /', $sql, $m)) {
            throw $this->rejected($sql); // phpcs:ignore WordPress.Security.EscapeOutput.ExceptionNotEscaped -- raw SQL; rejected() shortens and escapes it in the message
        }
        $t = $this->known($m[1], $sql);
        $clause = '/\GADD CONSTRAINT (`(?:[^`]|``)+`) FOREIGN KEY \((' . self::COLS . ')\) REFERENCES ' . self::IDENT . ' \((' . self::COLS . ')\)(' . self::ACTIONS . ')(, |$)/';
        $p = strlen($m[0]);
        $clauses = [];
        while ($p < strlen($sql)) {
            if (!preg_match($clause, $sql, $c, 0, $p)) {
                throw $this->rejected($sql); // phpcs:ignore WordPress.Security.EscapeOutput.ExceptionNotEscaped -- raw SQL; rejected() shortens and escapes it in the message
            }
            $ref = $this->known($c[3], $sql);
            $clauses[] = 'ADD CONSTRAINT ' . $c[1] . ' FOREIGN KEY (' . $c[2] . ') REFERENCES ' . Db::ident($this->final[$ref]) . ' (' . $c[4] . ')' . $c[5];
            $p += strlen($c[0]);
            if ($c[6] === '') {
                break;
            }
        }
        if ($clauses === [] || $p !== strlen($sql)) {
            throw $this->rejected($sql); // phpcs:ignore WordPress.Security.EscapeOutput.ExceptionNotEscaped -- raw SQL; rejected() shortens and escapes it in the message
        }
        return ['kind' => 'constraint', 'table' => $this->final[$t], 'sql' => 'ALTER TABLE ' . Db::ident($this->final[$t]) . ' ' . implode(', ', $clauses)];
    }

    private function known(string $quoted, string $sql): string
    {
        $t = str_replace('``', '`', $quoted);
        if (!isset($this->import[$t], $this->final[$t])) {
            throw $this->rejected($sql); // phpcs:ignore WordPress.Security.EscapeOutput.ExceptionNotEscaped -- raw SQL; rejected() shortens and escapes it in the message
        }
        return $t;
    }

    private function rejected(string $sql): RelayError
    {
        if ($this->from_site) {
            return new RelayError(
                'restore.statement_rejected',
                /* translators: %s: start of the rejected SQL statement */
                esc_html(sprintf(__('The other site sent SQL SiteRelay will not run: %s', 'siterelay'), substr($sql, 0, 100))),
                esc_html__('SiteRelay copies tables in the forms WordPress and its plugins create. Leave this table out (choose the tables under Custom), and tell SiteRelay’s developers which plugin made it.', 'siterelay')
            );
        }
        return new RelayError(
            'restore.statement_rejected',
            /* translators: %s: start of the rejected SQL statement */
            esc_html(sprintf(__('The backup contains SQL SiteRelay will not run: %s', 'siterelay'), substr($sql, 0, 100))),
            esc_html__('Only backups made by SiteRelay can be restored.', 'siterelay')
        );
    }
}