-
Notifications
You must be signed in to change notification settings - Fork 0
/
DBOperator.py
executable file
·96 lines (78 loc) · 2.92 KB
/
DBOperator.py
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
#!/usr/bin/python
from ftplib import FTP
import socket
import os
import pymongo
import ConfigParser
from tools import set_logger
import traceback
config_file = os.getcwd() + '/config.ini'
cf = ConfigParser.ConfigParser()
cf.read(config_file)
DB_ADDR = cf.get('mongodb','addr')
DB_PORT = int(cf.get('mongodb','port'),10)
READWRITE_USER = cf.get('mongodb','readwrite_user')
READWRITE_PASSWD = cf.get('mongodb','readwrite_password')
LOG_FILE = os.path.join(os.getcwd(),cf.get('support','log_path'))
logger = set_logger('DBOperator.py',LOG_FILE)
class DBOperator:
def __init__(self):
try:
self.connection = pymongo.Connection(DB_ADDR,DB_PORT)
#Database
self.db = self.connection.ftp
# logger.debug('Connected to ftp')
self.db.authenticate(READWRITE_USER,READWRITE_PASSWD)
# logger.debug('Database authenticate success!')
except:
# print 'Database connect error!'
traceback.print_exc()
logger.error('Database connect error,exit.')
exit(0)
def insert_to_allocate_queue(self,ip,client):
#collection
allocate_queue = (self.db).allocate_queue
#record = '{\"ip\":\"' + ip +'\",' + '\"client\":\"' + client + '\"}'
allocate_queue.insert({"_id":ip,"client":client})
def insert_ips_to_allocate(self,ips,client):
allocate_queue = self.db.allocate_queue
for i in ips:
allocate_queue.insert({"_id":i,"client":client})
def remove_from_allocate_queue(self,ip):
allocate_queue = self.db.allocate_queue
allocate_queue.remove({"_id":ip})
def insert_to_scanned_queue(self,ip,info,time,client):
scanned_queue = self.db.scanned_queue
scanned_queue.insert({"_id":ip,"info":info,"time":time,"client":client})
def clear_allocate_queue(self):
allocate_queue = self.db.allocate_queue
allocate_queue.remove()
def clear_scanned_queue(self):
scanned_queue = self.db.scanned_queue
scanned_queue.remove()
def get_allocate_queue_record_amount(self):
allocate_queue = self.db.allocate_queue
return allocate_queue.find().count()
def get_client_allocated_amount(self,client):
allocate_queue = self.db.allocate_queue
return allocate_queue.find({'client':client}).count()
def reallocate_ip(self,ip,to_client):
allocate_queue = self.db.allocate_queue
allocate_queue.update({'_id':ip},{'$set':{'client':to_client}})
def reallocate_ips(self,src_client,to_client):
allocate_queue = self.db.allocate_queue
allocate_queue.update({'client':src_client},{'$set':{'client':to_client}},upsert=False,multi=True)
def get_client_ips(self,client):
ip_list = []
allocate_queue = self.db.allocate_queue
for item in allocate_queue.find({'client':client}):
ip_list.append(item['_id'])
return ip_list
if __name__=="__main__":
test = DBOperator()
# test.insert_to_allocate_queue('127.0.0.1','client1')
# test.reallocate_ip('127.0.0.1','client2')
# test.reallocate_ips('client2','client1')
# print test.get_allocate_queue_record_amount()
# test.clear_allocate_queue()
test.remove_from_allocate_queue('192.168.128.99')