'''
Faraday Penetration Test IDE
Copyright (C) 2013  Infobyte LLC (http://www.infobytesec.com/)
See the file 'doc/LICENSE' for the license information

'''

import pytest

from faraday.server.utils.database import get_unique_fields
from faraday.server.models import (
    License,
    Service,
    Host,
    Vulnerability,
    Workspace,
    vulnerability_uniqueness
)

UNIQUE_FIELDS = {
    License: ['product', 'start_date', 'end_date'],
    Service: ['port', 'protocol', 'host_id', 'workspace_id'],
    Host: ['ip', 'workspace_id'],
    Vulnerability: [
        'name',
        'description',
        'type',
        'host_id',
        'service_id',
        'method',
        'parameter_name',
        'path',
        'website',
        'workspace_id',
    ],
    Workspace: ['name']
}


def test_vulnerability_ddl_invariant(session):
    """
        This test is to make sure that get_unique_fields for Vulnerability is
        returning the same columns as in the DDL
    :return:
    """
    statement = vulnerability_uniqueness.statement
    column_part = statement.split('%(fullname)s')[1]
    statements_clean = column_part.strip().strip(')').strip('(').split(',')
    statements_clean = [column.replace('COALESCE(', '').replace('md5(', '').strip('(') for column in statements_clean]
    statements_clean = \
        list(
            filter(len,
                   map(lambda column: column.strip("'')").strip('-1').strip('-1));').strip(), statements_clean)
                   )
        )
    statements_clean.remove('source_code_id')  # we don't support source_code yet
    unique_constraints = get_unique_fields(session, Vulnerability())
    for unique_constraint in unique_constraints:
        assert len(statements_clean) == len(unique_constraint)
        for statement_clean in statements_clean:
            statement_clean = statement_clean
            if statement_clean not in unique_constraint:
                raise Exception('Please check server.utils.database.get_unique_fields. Vulnerability DDL changed?')


@pytest.mark.parametrize("obj_class, expected_unique_fields", list(UNIQUE_FIELDS.items()))
def test_unique_fields_workspace(obj_class, expected_unique_fields, session):
    object_ = obj_class()
    unique_constraints = get_unique_fields(session, object_)
    for unique_constraint in unique_constraints:
        assert unique_constraint == expected_unique_fields