#!/usr/bin/env python

"""
Copyright (c) 2006-2026 sqlmap developers (https://sqlmap.org)
See the file 'LICENSE' for copying permission
"""

import re

from lib.core.agent import agent
from lib.core.bigarray import BigArray
from lib.core.common import Backend
from lib.core.common import isNoneValue
from lib.core.common import singleTimeWarnMessage
from lib.core.common import unArrayizeValue
from lib.core.common import unsafeSQLIdentificatorNaming
from lib.core.compat import xrange
from lib.core.convert import getConsoleLength
from lib.core.convert import getUnicode
from lib.core.data import conf
from lib.core.data import logger
from lib.core.data import queries
from lib.core.dicts import DUMP_REPLACEMENTS
from lib.core.enums import CHARSET_TYPE
from lib.core.enums import DBMS
from lib.core.enums import EXPECTED
from lib.core.settings import NULL
from lib.core.unescaper import unescaper
from lib.request import inject
from lib.utils.safe2bin import safechardecode

# back-end DBMSes whose dump table reference is schema/database-qualified (db.table).
# Note: for MSSQL the table identifier already carries its schema (e.g. dbo.users), so the
# plain db.table form yields the correct db.schema.table (e.g. [master].dbo.users).
KEYSET_SCHEMA_QUALIFIED = (DBMS.MYSQL, DBMS.PGSQL, DBMS.CRATEDB, DBMS.MSSQL, DBMS.H2, DBMS.HSQLDB)

def _tableRef(tbl):
    dbms = Backend.getIdentifiedDbms()
    if dbms in (DBMS.ORACLE,) and conf.db:
        return "%s.%s" % (conf.db.upper(), tbl.upper())
    if dbms in KEYSET_SCHEMA_QUALIFIED and conf.db:
        return "%s.%s" % (conf.db, tbl)
    return tbl

def keysetSupported():
    """
    Whether the back-end DBMS declares the keyset (seek) pagination queries and a
    cursor source (a physical row-id pseudo-column or a primary-key catalog lookup)
    """

    dumpNode = queries[Backend.getIdentifiedDbms()].dump_table
    return "keyset_next" in dumpNode.blind and ("rowid" in dumpNode.blind or "primary_key" in dumpNode)

def _integerCursor(tbl, cursor):
    """
    Whether every cursor column holds integer values, probed via MIN(col).

    Only integer keys are accepted: _embed() emits them as bare numeric literals, giving a
    numeric comparison that matches MIN/ORDER BY. String (and even decimal) keys would be
    escaped to a binary/hex literal whose order can differ from MIN's collation and silently
    skip rows, so they are rejected here and fall back to the OFFSET dump.
    """

    blind = queries[Backend.getIdentifiedDbms()].dump_table.blind
    ref = _tableRef(tbl)

    for column in cursor:
        query = agent.whereQuery(blind.keyset_first % (agent.preprocessField(tbl, column), ref))
        value = unArrayizeValue(inject.getValue(query))

        # empty/NULL MIN (e.g. empty table) is not disqualifying; the walk just yields no rows
        if not isNoneValue(value) and re.match(r"\A-?[0-9]+\Z", getUnicode(value).strip()) is None:
            return False

    return True

def resolveKeysetCursor(tbl, colList):
    """
    Returns the list of column(s) forming a stable, indexed cursor for keyset (seek)
    pagination of the table: a declared physical row-id pseudo-column when available,
    otherwise the indexed primary key (single or composite) resolved from the catalog.
    Returns None when neither applies or a key column is not part of the dumped columns.
    """

    if not keysetSupported():
        return None

    dumpNode = queries[Backend.getIdentifiedDbms()].dump_table

    # 1) a declared physical row-id pseudo-column (always unique + indexed where supported)
    if "rowid" in dumpNode.blind:
        return [dumpNode.blind.rowid]

    # 2) the indexed primary key (single-column, or composite when keyset_ordered is declared)
    pkNode = dumpNode.primary_key

    # Note: schema/table are string literals in the catalog lookups, so the unquoted
    # (identifier-unescaped) names are used (the dump queries keep the quoted form)
    unsafeDb = unsafeSQLIdentificatorNaming(conf.db)
    unsafeTbl = unsafeSQLIdentificatorNaming(tbl)

    # Note: no whereQuery() here - these are catalog (schema) lookups, so the data-row
    # filter from --where must not be appended to them
    query = pkNode.count % (unsafeDb, unsafeTbl)
    count = inject.getValue(query, expected=EXPECTED.INT, charsetType=CHARSET_TYPE.DIGITS)

    try:
        count = int(count)
    except (ValueError, TypeError):
        return None

    if count < 1:
        return None

    # composite keys require the row-value/ordered keyset form
    if count > 1 and "keyset_ordered" not in dumpNode.blind:
        return None

    cursor = []
    for index in xrange(count):
        query = pkNode.query % (unsafeDb, unsafeTbl, index)
        column = unArrayizeValue(inject.getValue(query))

        if not column:
            return None

        match = None
        for _ in colList:
            if _ and _.lower() == column.lower():
                match = _
                break

        if match is None:
            return None

        cursor.append(match)

    # restrict to integer cursors: a string key's escaped-literal comparison may order
    # differently than MIN/ORDER BY and silently skip rows (such keys fall back to OFFSET)
    if not _integerCursor(tbl, cursor):
        return None

    return cursor

def _lit(value):
    """
    Type-correct SQL literal for a cursor value: a bare numeric literal for numeric keys
    (so the index is still used and the comparison is numeric), otherwise the DBMS-escaped
    (e.g. 0x.. hex) form for string keys. Both forms are self-contained (no surrounding quotes).
    """

    if value is None:
        return NULL                                 # unescaper.escape() passes None through, and a bare
                                                    # None formatted into a predicate is not even SQL
    if re.match(r"\A-?[0-9]+\Z", value):
        return value
    return unescaper.escape(value, False)

def _embed(template, value, *fixed):
    """
    Fills a single-column keyset template whose trailing placeholder is the cursor value.
    """

    template = template.replace("'%s'", "%s")
    return template % (fixed + (_lit(value),))

def _target(count):
    """Rows the walk is expected to produce, honouring --start/--stop."""

    if conf.limitStart and conf.limitStop:
        return max(0, conf.limitStop - conf.limitStart + 1)
    elif conf.limitStop:
        return conf.limitStop
    elif conf.limitStart:
        return max(0, count - conf.limitStart + 1)

    return count

def _dumpSingle(tbl, colList, count, cursor, tableRef, entries, lengths):
    """False when the walk gave up mid-table (the caller then discards the partial result)."""

    blind = queries[Backend.getIdentifiedDbms()].dump_table.blind
    field = agent.preprocessField(tbl, cursor)
    target = _target(count)

    pivotValue = None

    # hybrid: a single OFFSET jump to seed the cursor just before --start, then pure keyset
    if conf.limitStart and conf.limitStart > 1 and "keyset_seed" in blind:
        query = agent.whereQuery(blind.keyset_seed % (field, tableRef, field, conf.limitStart - 2))
        seed = unArrayizeValue(inject.getValue(query))

        if isNoneValue(seed) or seed == NULL:
            return False                            # no seed, no walk - and an empty table is not that

        pivotValue = safechardecode(seed)

    produced = 0

    while produced < target:
        # Advance with ORDER BY ... LIMIT 1 (like the composite path), NOT MIN(): the value-extraction
        # casts the aggregated column to VARCHAR *inside* MIN(), yielding a LEXICAL minimum ('10' after
        # '1') that disagrees with the numeric '>' comparison and silently skips rows (2..9, 11..). The
        # ORDER BY is on the raw (numeric) column, so the next cursor value is the true successor.
        condition = "1=1" if pivotValue is None else "%s>%s" % (field, _lit(pivotValue))
        query = agent.whereQuery(blind.keyset_ordered % (field, tableRef, condition, field))
        value = unArrayizeValue(inject.getValue(query))

        if isNoneValue(value) or value == NULL:
            break

        value = safechardecode(value)

        # safety latch against a non-advancing cursor (e.g. encoding edge cases)
        if value == pivotValue:
            singleTimeWarnMessage("keyset cursor stopped advancing prematurely")
            return False

        pivotValue = value

        for column in colList:
            if column == cursor:
                colValue = pivotValue
            else:
                query = _embed(blind.keyset_by, pivotValue, agent.preprocessField(tbl, column), tableRef, field)
                query = agent.whereQuery(query)
                colValue = unArrayizeValue(inject.getValue(query, dump=True))

            colValue = "" if isNoneValue(colValue) else colValue
            lengths[column] = max(lengths[column], getConsoleLength(DUMP_REPLACEMENTS.get(getUnicode(colValue), getUnicode(colValue))))
            entries[column].append(colValue)

        produced += 1

    return True

def _dumpComposite(tbl, colList, count, cursorCols, tableRef, entries, lengths):
    """False when the walk gave up mid-table (the caller then discards the partial result)."""

    blind = queries[Backend.getIdentifiedDbms()].dump_table.blind
    fields = [agent.preprocessField(tbl, _) for _ in cursorCols]
    orderExpr = ','.join(fields)

    startSkip = (conf.limitStart - 1) if conf.limitStart else 0
    target = _target(count)

    prev = None
    produced = 0
    seen = 0

    while produced < target and seen < count:
        if prev is None:
            condition = "1=1"
        else:
            # Portable lexicographic seek predicate. ANSI row-value comparison ((a,b)>(x,y)) is not
            # supported on MSSQL/Oracle - there it errored, stopping the walk after the first row -
            # so expand it to (a>x) OR (a=x AND b>y) OR ... which uses only scalar comparisons.
            ors = []
            for i in xrange(len(fields)):
                terms = ["%s=%s" % (fields[j], _lit(prev[j])) for j in xrange(i)]
                terms.append("%s>%s" % (fields[i], _lit(prev[i])))
                ors.append("(%s)" % " AND ".join(terms))
            condition = "(%s)" % " OR ".join(ors)

        tup = []
        for field in fields:
            query = agent.whereQuery(blind.keyset_ordered % (field, tableRef, condition, orderExpr))
            value = unArrayizeValue(inject.getValue(query))
            tup.append(None if isNoneValue(value) else safechardecode(value))

        if all(isNoneValue(_) for _ in tup):
            break                                   # nothing past the cursor: the table is walked

        # A key column that did not come back (an error-channel miss, a blocked payload) cannot be
        # seeked on, and its equality would pin the rest of the row to a NULL - so the walk stops
        # here rather than emitting a row of empty cells and carrying the hole into the next seek
        if any(isNoneValue(_) for _ in tup):
            singleTimeWarnMessage("keyset cursor could not be retrieved for one of the key column(s)")
            return False

        if prev is not None and tup == prev:
            singleTimeWarnMessage("keyset cursor stopped advancing prematurely")
            return False

        prev = tup
        seen += 1

        if seen <= startSkip:
            continue

        equals = " AND ".join("%s=%s" % (field, _lit(value)) for field, value in zip(fields, tup))

        for column in colList:
            if column in cursorCols:
                colValue = tup[cursorCols.index(column)]
            else:
                query = agent.whereQuery(blind.keyset_where % (agent.preprocessField(tbl, column), tableRef, equals))
                colValue = unArrayizeValue(inject.getValue(query, dump=True))

            colValue = "" if isNoneValue(colValue) else colValue
            lengths[column] = max(lengths[column], getConsoleLength(DUMP_REPLACEMENTS.get(getUnicode(colValue), getUnicode(colValue))))
            entries[column].append(colValue)

        produced += 1

    return True

def keysetDumpTable(tbl, colList, count, cursor):
    """
    Dumps a table one row at a time using keyset (seek) pagination on 'cursor' (a list of
    one or more indexed key columns): the next row is reached with a >/row-value comparison
    against the previous cursor (index range scan) and every other column is fetched with an
    exact equality on the cursor (index point seek), so no row is skipped via OFFSET and no
    per-row ORDER BY filesort is needed. A deep --start uses a single OFFSET "seed" jump
    (single-column cursors), after which the walk is pure keyset.

    Returns None when the walk gave up mid-table (a key value that did not come back, a cursor
    that stopped advancing): a short table is worse than a slow one, so the caller redoes it
    with the standard OFFSET dump instead of showing whatever was reached.
    """

    tableRef = _tableRef(tbl)
    lengths = {}
    entries = {}

    for column in colList:
        lengths[column] = 0
        entries[column] = BigArray()

    if len(cursor) == 1:
        complete = _dumpSingle(tbl, colList, count, cursor[0], tableRef, entries, lengths)
    else:
        complete = _dumpComposite(tbl, colList, count, cursor, tableRef, entries, lengths)

    if not complete:
        warnMsg = "keyset pagination did not complete for table '%s', " % unsafeSQLIdentificatorNaming(tbl)
        warnMsg += "falling back to the standard dump"
        logger.warning(warnMsg)

        return None

    debugMsg = "keyset pagination retrieved %d row(s) for table '%s'" % (len(entries[colList[0]]) if colList and colList[0] in entries else 0, unsafeSQLIdentificatorNaming(tbl))
    logger.debug(debugMsg)

    return entries, lengths