JSShell/db_handler.py at master · ezrac/JSShell · 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
from app import app, db
from app.models import Client, Command
import argparse
import uuid

def sqltry(func):
    def wrap(*args, **kwargs):
        try:
            func(*args, **kwargs)
        except Exception as e:
            print('Error occurred:', e)
            return None
    return wrap


@sqltry
def createTable():
    print('Creating table...')
    db.create_all()
    print('Done!')


@sqltry
def listRecords():
    print('Listing records...')
    clients = Client.query.all()

    for c in clients:
        print(c)


@sqltry
def insertRecord(_idDefault = str(uuid.uuid4()), userAgentDefault = 'Mozilla Testing/1.0', ipDefault = '127.0.0.1'):
    print('Enter Values: ')

    _id = input('UUID: ')
    if not _id: _id = _idDefault

    user_agent = input('User-Agent: ')
    if not user_agent: user_agent = userAgentDefault

    ip = input('IP: ')
    if not ip: ip = ipDefault

    c = Client(_id, user_agent, ip)
    db.session.add(c)
    db.session.commit()


@sqltry
def insertDummy():
    _id = str(uuid.uuid4())
    user_agent = 'Mozilla Testing/1.0'
    ip = '127.0.0.1'
    c = Client(_id, user_agent, ip)
    db.session.add(c)
    db.session.commit()


@sqltry
def dropTable():
    if input('Sure? [y/n]') == 'y':
        db.drop_all()
        print('Table dropped.')
    else:
        print('Bad choice. Bye!')

@sqltry
def dropCreateList():
    dropTable()
    createTable()
    listRecords()


@sqltry
def truncateTable():
    Client.query.delete()
    db.session.commit()


@sqltry
def createCommand():
    i = input('Enter Client ID: ')
    c = Command('aaa','bbb')
    u = Client.query.filter_by(id=int(i)).first()
    u.commands.append(c)
    db.session.add(c)
    db.session.commit()


if __name__ == "__main__":

    actions = {
        'list' : listRecords,
        'create' : createTable,
        'insert' : insertRecord,
        'dummy' : insertDummy,
        'drop' : dropTable,
        'dcl' : dropCreateList,
        'trunc' : truncateTable,
        'com' : createCommand
    }

    parser = argparse.ArgumentParser(description='DB Handler')
    parser.add_argument('action', choices=actions.keys())


    args = parser.parse_args()
    actions[args.action]()