Codebase list python-faraday / 62d1b14 faraday / migrations / versions / 15d70093d262_severities_histogram_model.py
62d1b14

Tree @62d1b14 (Download .tar.gz)

15d70093d262_severities_histogram_model.py @62d1b14 — raw · history · blame 

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
"""Severities histogram model

Revision ID: 15d70093d262
Revises: d8f0b32a5c0e
Create Date: 2021-11-08 13:57:28.099487+00:00

"""
from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
from sqlalchemy import func, case

from faraday.server.models import VulnerabilityGeneric, SeveritiesHistogram, Workspace

revision = '15d70093d262'
down_revision = 'd8f0b32a5c0e'
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('severities_histogram',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('workspace_id', sa.Integer(), nullable=False),
    sa.Column('date', sa.Date(), nullable=False),
    sa.Column('medium', sa.Integer(), nullable=False),
    sa.Column('high', sa.Integer(), nullable=False),
    sa.Column('critical', sa.Integer(), nullable=False),
    sa.Column('confirmed', sa.Integer(), nullable=False),
    sa.ForeignKeyConstraint(['workspace_id'], ['workspace.id'], ),
    sa.PrimaryKeyConstraint('id')
    )
    op.create_index(op.f('ix_severities_histogram_workspace_id'), 'severities_histogram', ['workspace_id'], unique=False)
    # ### end Alembic commands ###

    # Init histogram
    bind = op.get_bind()
    session = sa.orm.Session(bind=bind)
    workspaces = session.query(Workspace).all()
    for workspace in workspaces:
        vulnerabilities = session.query(VulnerabilityGeneric) \
            .with_entities(func.date_trunc('day', VulnerabilityGeneric.create_date),
                           VulnerabilityGeneric.severity,
                           func.count(VulnerabilityGeneric.severity),
                           func.sum(case([(VulnerabilityGeneric.confirmed, 1)], else_=0)))\
            .filter(VulnerabilityGeneric.workspace_id == workspace.id,
                    VulnerabilityGeneric.status.notin_(['closed', 'risk-accepted']),
                    VulnerabilityGeneric.severity.in_(['medium', 'high', 'critical']))\
            .group_by(func.date_trunc('day', VulnerabilityGeneric.create_date), VulnerabilityGeneric.severity).all()
        for histogram_date, severity_type, severity_count, confirmed_count in vulnerabilities:
            severity_histogram = session.query(SeveritiesHistogram)\
                .filter(SeveritiesHistogram.date == histogram_date,
                        SeveritiesHistogram.workspace_id == workspace.id).first()
            if severity_histogram is None:
                severity_histogram = SeveritiesHistogram(date=histogram_date, workspace=workspace, medium=0, high=0, critical=0, confirmed=0)
                session.add(severity_histogram)
                session.commit()
            if severity_type == 'medium':
                severity_histogram.medium = severity_count
            if severity_type == 'high':
                severity_histogram.high = severity_count
            if severity_type == 'critical':
                severity_histogram.critical = severity_count
            severity_histogram.confirmed += confirmed_count
        session.commit()


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index(op.f('ix_severities_histogram_workspace_id'), table_name='severities_histogram')
    op.drop_table('severities_histogram')
    # ### end Alembic commands ###