관리-도구

편집 파일: 070_modsec_incident_names.py

from tempfile import TemporaryFile

def extract_name(description):
    return description.split("||", maxsplit=1)[0]

def migrate(migrator, database, fake=False, **kwargs):
    """
    This migration extracts incident name from whole mod_security message
    Centos6 version of sqlite does not have instr(), using slow python-based
    way
    """
    incident = migrator.orm["incident"]

# FIXME: after migrating to peewee 3 remove the try..except block
    def select_incidents():
        try:
            yield from (
                incident.select(incident.id, incident.description)
                .where(incident.plugin == "modsec")
                .where(incident.description.contains("||"))
                .tuples()
                .iterator()
            )
        except RuntimeError:
            return

with TemporaryFile(mode="w+") as f:
        for id_, desc in select_incidents():
            f.write("{},{}\n".format(id_, extract_name(desc)))
        f.seek(0)
        with database.atomic():
            for line in f:
                id_, name = line.split(",", maxsplit=1)
                incident.update(name=name.strip()).where(
                    incident.id == int(id_)
                ).execute()

def rollback(migrator, database, fake=False, **kwargs):
    """Write your rollback migrations here."""
    pass