Skip to content

Commit 08ab8d4

Browse files
xmunozewdurbin
xmunoz
authored and
ewdurbin
committed
Add new models for malware detection. (#7118)
* Add new models for malware detection. Fixes #7090 and #7092. * Code review changes. - FK on release_file.id field instead of md5 - Change message type from String to Text - Change Enum class in model to singular form
1 parent 9f134c4 commit 08ab8d4

File tree

3 files changed

+237
-0
lines changed

3 files changed

+237
-0
lines changed

‎warehouse/malware/__init__.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
# Licensed under the Apache License, Version 2.0 (the "License");
2+
# you may not use this file except in compliance with the License.
3+
# You may obtain a copy of the License at
4+
#
5+
# http://www.apache.org/licenses/LICENSE-2.0
6+
#
7+
# Unless required by applicable law or agreed to in writing, software
8+
# distributed under the License is distributed on an "AS IS" BASIS,
9+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10+
# See the License for the specific language governing permissions and
11+
# limitations under the License.

‎warehouse/malware/models.py

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
# Licensed under the Apache License, Version 2.0 (the "License");
2+
# you may not use this file except in compliance with the License.
3+
# You may obtain a copy of the License at
4+
#
5+
# http://www.apache.org/licenses/LICENSE-2.0
6+
#
7+
# Unless required by applicable law or agreed to in writing, software
8+
# distributed under the License is distributed on an "AS IS" BASIS,
9+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10+
# See the License for the specific language governing permissions and
11+
# limitations under the License.
12+
13+
import enum
14+
15+
from citext import CIText
16+
from sqlalchemy import (
17+
Boolean,
18+
Column,
19+
DateTime,
20+
Enum,
21+
ForeignKey,
22+
Integer,
23+
String,
24+
Text,
25+
UniqueConstraint,
26+
orm,
27+
sql,
28+
)
29+
from sqlalchemy.dialects.postgresql import JSONB
30+
31+
from warehouse import db
32+
from warehouse.utils.attrs import make_repr
33+
34+
35+
class MalwareCheckType(enum.Enum):
36+
37+
EventHook = "event_hook"
38+
Scheduled = "scheduled"
39+
40+
41+
class MalwareCheckState(enum.Enum):
42+
43+
Enabled = "enabled"
44+
Evaluation = "evaluation"
45+
Disabled = "disabled"
46+
WipedOut = "wiped_out"
47+
48+
49+
class VerdictClassification(enum.Enum):
50+
51+
Threat = "threat"
52+
Indeterminate = "indeterminate"
53+
Benign = "benign"
54+
55+
56+
class VerdictConfidence(enum.Enum):
57+
58+
Low = "low"
59+
Medium = "medium"
60+
High = "high"
61+
62+
63+
class MalwareCheck(db.Model):
64+
65+
__tablename__ = "malware_checks"
66+
__table_args__ = (UniqueConstraint("name", "version"),)
67+
__repr__ = make_repr("name", "version")
68+
69+
name = Column(CIText, nullable=False)
70+
version = Column(Integer, default=0, nullable=False)
71+
short_description = Column(String(length=128), nullable=False)
72+
long_description = Column(Text, nullable=False)
73+
check_type = Column(
74+
Enum(MalwareCheckType, values_callable=lambda x: [e.value for e in x]),
75+
nullable=False,
76+
)
77+
# This field contains the same content as the ProjectEvent and UserEvent "tag"
78+
# fields.
79+
hook_name = Column(String, nullable=True)
80+
state = Column(
81+
Enum(MalwareCheckState, values_callable=lambda x: [e.value for e in x]),
82+
nullable=False,
83+
server_default=("disabled"),
84+
)
85+
created = Column(DateTime, nullable=False, server_default=sql.func.now())
86+
87+
88+
class MalwareVerdict(db.Model):
89+
__tablename__ = "malware_verdicts"
90+
91+
run_date = Column(DateTime, nullable=False, server_default=sql.func.now())
92+
check_id = Column(
93+
ForeignKey("malware_checks.id", onupdate="CASCADE", ondelete="CASCADE"),
94+
nullable=False,
95+
index=True,
96+
)
97+
file_id = Column(ForeignKey("release_files.id"), nullable=False)
98+
classification = Column(
99+
Enum(VerdictClassification, values_callable=lambda x: [e.value for e in x]),
100+
nullable=False,
101+
)
102+
confidence = Column(
103+
Enum(VerdictConfidence, values_callable=lambda x: [e.value for e in x]),
104+
nullable=False,
105+
)
106+
message = Column(Text, nullable=True)
107+
details = Column(JSONB, nullable=True)
108+
manually_reviewed = Column(Boolean, nullable=False, server_default=sql.false())
109+
administrator_verdict = Column(
110+
Enum(VerdictClassification, values_callable=lambda x: [e.value for e in x]),
111+
nullable=True,
112+
)
113+
full_report_link = Column(String, nullable=True)
114+
115+
check = orm.relationship("MalwareCheck", foreign_keys=[check_id], lazy=True)
116+
release_file = orm.relationship("File", foreign_keys=[file_id], lazy=True)

‎warehouse/migrations/versions/061ff3d24c22_add_malware_detection_tables.py

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
# Licensed under the Apache License, Version 2.0 (the "License");
2+
# you may not use this file except in compliance with the License.
3+
# You may obtain a copy of the License at
4+
#
5+
# http://www.apache.org/licenses/LICENSE-2.0
6+
#
7+
# Unless required by applicable law or agreed to in writing, software
8+
# distributed under the License is distributed on an "AS IS" BASIS,
9+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10+
# See the License for the specific language governing permissions and
11+
# limitations under the License.
12+
"""
13+
Add malware detection tables
14+
15+
Revision ID: 061ff3d24c22
16+
Revises: b5bb5d08543d
17+
Create Date: 2019-12-18 17:27:00.183542
18+
"""
19+
import citext
20+
import sqlalchemy as sa
21+
22+
from alembic import op
23+
from sqlalchemy.dialects import postgresql
24+
25+
revision = "061ff3d24c22"
26+
down_revision = "b5bb5d08543d"
27+
28+
MalwareCheckTypes = sa.Enum("event_hook", "scheduled", name="malwarechecktypes")
29+
30+
MalwareCheckStates = sa.Enum(
31+
"enabled", "evaluation", "disabled", "wiped_out", name="malwarecheckstate"
32+
)
33+
34+
VerdictClassifications = sa.Enum(
35+
"threat", "indeterminate", "benign", name="verdictclassification"
36+
)
37+
VerdictConfidences = sa.Enum("low", "medium", "high", name="verdictconfidence")
38+
39+
40+
def upgrade():
41+
op.create_table(
42+
"malware_checks",
43+
sa.Column(
44+
"id",
45+
postgresql.UUID(as_uuid=True),
46+
server_default=sa.text("gen_random_uuid()"),
47+
nullable=False,
48+
),
49+
sa.Column("name", citext.CIText(), nullable=False),
50+
sa.Column("version", sa.Integer(), nullable=False),
51+
sa.Column("short_description", sa.String(length=128), nullable=False),
52+
sa.Column("long_description", sa.Text(), nullable=False),
53+
sa.Column("check_type", MalwareCheckTypes, nullable=False),
54+
sa.Column("hook_name", sa.String(), nullable=True),
55+
sa.Column(
56+
"state", MalwareCheckStates, server_default="disabled", nullable=False,
57+
),
58+
sa.Column(
59+
"created", sa.DateTime(), server_default=sa.text("now()"), nullable=False
60+
),
61+
sa.PrimaryKeyConstraint("id"),
62+
sa.UniqueConstraint("name", "version"),
63+
)
64+
op.create_table(
65+
"malware_verdicts",
66+
sa.Column(
67+
"id",
68+
postgresql.UUID(as_uuid=True),
69+
server_default=sa.text("gen_random_uuid()"),
70+
nullable=False,
71+
),
72+
sa.Column(
73+
"run_date", sa.DateTime(), server_default=sa.text("now()"), nullable=False
74+
),
75+
sa.Column("check_id", postgresql.UUID(as_uuid=True), nullable=False),
76+
sa.Column("file_id", postgresql.UUID(as_uuid=True), nullable=False),
77+
sa.Column("classification", VerdictClassifications, nullable=False,),
78+
sa.Column("confidence", VerdictConfidences, nullable=False,),
79+
sa.Column("message", sa.Text(), nullable=True),
80+
sa.Column("details", postgresql.JSONB(astext_type=sa.Text()), nullable=True),
81+
sa.Column(
82+
"manually_reviewed",
83+
sa.Boolean(),
84+
server_default=sa.text("false"),
85+
nullable=False,
86+
),
87+
sa.Column("administrator_verdict", VerdictClassifications, nullable=True,),
88+
sa.Column("full_report_link", sa.String(), nullable=True),
89+
sa.ForeignKeyConstraint(
90+
["check_id"], ["malware_checks.id"], onupdate="CASCADE", ondelete="CASCADE"
91+
),
92+
sa.ForeignKeyConstraint(["file_id"], ["release_files.id"]),
93+
sa.PrimaryKeyConstraint("id"),
94+
)
95+
op.create_index(
96+
op.f("ix_malware_verdicts_check_id"),
97+
"malware_verdicts",
98+
["check_id"],
99+
unique=False,
100+
)
101+
102+
103+
def downgrade():
104+
op.drop_index(op.f("ix_malware_verdicts_check_id"), table_name="malware_verdicts")
105+
op.drop_table("malware_verdicts")
106+
op.drop_table("malware_checks")
107+
MalwareCheckTypes.drop(op.get_bind())
108+
MalwareCheckStates.drop(op.get_bind())
109+
VerdictClassifications.drop(op.get_bind())
110+
VerdictConfidences.drop(op.get_bind())

0 commit comments

Comments
 (0)