PGPool/pgpool-import.py at master · codename-art/PGPool · 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
import codecs
import logging
import os

import sys
from flask import Flask

from pgpool.config import args
from pgpool.models import init_database, Account

logging.basicConfig(level=logging.INFO,
    format='%(asctime)s [%(threadName)16s][%(module)14s][%(levelname)8s] %(message)s')
log = logging.getLogger(__name__)

app = Flask(__name__)


def load_accounts_file(filename):
    accounts = []
    log.info("Loading accounts from file {}.".format(filename))
    with codecs.open(filename, mode='r', encoding='utf-8') as f:
        for line in f:
            if line.strip() == "":
                continue
            auth = usr = pwd = None
            fields = line.split(",")
            if len(fields) == 3:
                auth = fields[0].strip()
                usr = fields[1].strip()
                pwd = fields[2].strip()
            elif len(fields) == 2:
                auth = 'ptc'
                usr = fields[0].strip()
                pwd = fields[1].strip()
            elif len(fields) == 1:
                fields = line.split(":")
                auth = 'ptc'
                usr = fields[0].strip()
                pwd = fields[1].strip()
            if auth is not None:
                accounts.append({
                    'auth_service': auth,
                    'username': usr,
                    'password': pwd
                })
    if len(accounts) == 0:
        log.error("Could not load any accounts. Nothing to do. Exiting.")
        sys.exit(1)
    return accounts


def force_account_condition(account):
    account.ban_flag = 0
    if args.condition == 'good':
        account.banned = 0
        account.shadowbanned = 0
        account.captcha = 0
    elif args.condition == 'banned':
        account.banned = 1
        account.shadowbanned = 0
        account.captcha = 0
    elif args.condition == 'blind':
        account.banned = 0
        account.shadowbanned = 1
        account.captcha = 0
    elif args.condition == 'captcha':
        account.banned = 0
        account.shadowbanned = 0
        account.captcha = 1


# ---------------------------------------------------------------------------

log.info("PGPool CSV Importer starting up...")

db = init_database(app)

filename = args.import_csv
if not os.path.isfile(filename):
    log.error("File {} does not exist.".format(filename))
    sys.exit(1)

accounts = load_accounts_file(filename)
num_accounts = len(accounts)
log.info("Found {} accounts.".format(num_accounts))

num_skipped = 0
num_imported = 0

for acc in accounts:
    username = acc['username']
    account, created = Account.get_or_create(username=username)
    if created:
        account.auth_service = acc['auth_service']
        account.password = acc['password']
        account.level = args.level
        if args.condition != 'unknown':
            force_account_condition(account)
        account.save()
        addl_logmsg = ""
        if args.level:
            addl_logmsg += " | Forced trainer level: {}".format(args.level)
        addl_logmsg += " | Initial condition: {}".format(args.condition)
        log.info("Added account {}{}".format(username, addl_logmsg))
        num_imported += 1
    else:
        log.info("Account {} already known. Skipping.".format(username))
        num_skipped += 1

log.info("Done. Imported {} new accounts, skipped {} accounts.".format(num_imported, num_skipped))