exploitdbspider/db.py at master · wugedz/exploitdbspider · GitHub

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
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
# -*- coding: UTF-8 -*-
import sys

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy import Column, Integer, String

Base = declarative_base()


class DBPoc(Base):
    """
    create table exp_poc_info
    (
    id int(11) primary key,
    eid int(11),
    cve varchar(128),

    title varchar(512),
    author varchar(128),
    published_time varchar(128),

    verified varchar(128),

    platform varchar(128),
    exploit_type varchar(128),
    exploit_url varchar(128),
    exploit_app varchar(128)
    )
    """
    __tablename__ = 'exp_poc_info'

    id = Column(Integer, primary_key=True)
    eid = Column(Integer)
    cve = Column(String)

    title = Column(String)
    author = Column(String)
    published_time = Column(String)

    verified = Column(String)

    platform = Column(String)
    exploit_type = Column(String)
    exploit_url = Column(String)
    exploit_app = Column(String)

    def __init__(self, eid, cve,
                 title, author, published_time, verified,
                 platform, exploit_type, exploit_url, exploit_app):
        self.eid = eid
        self.cve = cve

        self.title = title
        self.author = author
        self.published_time = published_time
        self.verified = verified

        self.platform = platform
        self.exploit_type = exploit_type
        self.exploit_url = exploit_url
        self.exploit_app = exploit_app

    def __repr__(self):
        return "<id:%s poc_id ('%s')>:%s" % (self.id, self.eid, self.title)


class DBEngine(object):
    def __init__(self):
        #
        self.engine = create_engine('sqlite:///exploit_db.sqlite', echo=False)
        db_session = sessionmaker(autocommit=False, autoflush=False, bind=self.engine)
        self.session = db_session()

    def close_db(self):
        #
        self.session.close()

    # interface
    # lower words connected with '_'
    def add_poc(self, Poc):
        # 添加poc
        self.session.add(Poc)
        self.session.commit()

    def del_poc(self, eid):
        # 删除poc
        poc = self.session.query(DBPoc).filter(DBPoc.eid == eid).first()
        try:
            self.session.delete(poc)
            self.session.commit()
        except Exception as e:
            print(e)

    def is_eid_exist(self, eid):
        # exist True
        # not exist False
        if self.session.query(DBPoc).filter(DBPoc.eid == eid).first() is None:
            return False
        else:
            return True

    def view_all_poc(self):
        print('DBPoc:')
        all_poc = self.session.query(DBPoc)
        for poc in all_poc:
            print(poc)


if __name__ == '__main__':
    db = DBEngine()
    if len(sys.argv) >= 2 and sys.argv[1] == 'init':
        Base.metadata.create_all(db.engine)
    else:
        db.view_all_poc()