class SecBot:
def __init__(self, name, domain, room, login, password):
self.name = name
self.login = login
self.domain = domain
self.password = password
self.room = room
self.room_id = ''
self.api = None
self.processed = []
self.targets = {
'192.168.1.1': {
'comments': 'Router!',
'services': {
'21': 'ftp',
'22': 'ssh',
'9989': 'http',
}
},
'192.168.1.3': {
'comments': '',
'services': {
'80': 'http',
}
},
}
self.load_status()
self.api = RocketChatAPI(settings={
'username': self.login,
'password': self.password,
'domain': self.domain,
})
self.room_id = self.get_room_id('pentest-rdc')
def read_messages(self):
res = self.api.get_private_room_history(self.room_id)
for r in res['messages']:
yield r
def post(self, message):
self.api.send_message(message, self.room)
def process_message(self, message):
# {
# '_id': '9ota4D4ewkJpEmFBT',
# 'rid': 'TJf4aSR5Bo765XZk8',
# 'msg': 'Is Bob there?',
# 'ts': '2018-12-02T14:08:21.850Z',
# 'u': {
# '_id': 'Nd3f8bzCzSmWrLc32',
# 'username': '******',
# 'name': 'Tiao'
# },
# '_updatedAt': '2018-12-02T14:08:21.863Z',
# 'editedBy': None,
# 'editedAt': None,
# 'emoji': None,
# 'avatar': None,
# 'alias': None,
# 'customFields': None,
# 'groupable': None,
# 'attachments': None,
# 'reactions': None,
# 'mentions': [],
# 'channels': []
# }
# TODO: Process zenmap output
# TODO: Allow users to book one IP for investigation
# TODO: List non-investigated IPs
# TODO: List ports for one IP
# TODO: Propose tools commands to execute
if message['_id'] not in self.processed and message['u']['username'] != self.login:
print('Processing', message)
if 'Bob' in message['msg']:
s = message['u']['username']
self.post('@%s, who is Bob?' % s)
if '@secbot' in message['msg']:
s = message['u']['username']
if 'targets' in message['msg']:
print('listing targets')
msg = ''
ts = self.targets
for i in ts:
msg += i + ' ' + ts[i]['comments'] + '\n'
print('test')
self.post(msg)
self.processed.append(message['_id'])
def get_room_id(self, room):
rooms = self.api.get_private_rooms()
print(rooms)
for r in rooms:
if r['name'] == room:
return r['id']
return ''
def load_status(self):
with open('status.json', 'r') as f:
data = json.load(f)
self.processed = data['processed']
print('Processed', self.processed)
def save_status(self):
data = {}
data['processed'] = self.processed
with open('status.json', 'w') as f:
json.dump(data, f)
def start(self):
while 1:
for m in self.read_messages():
self.process_message(m)
time.sleep(0.5)
class Bot():
"""
创建一个聊天机器人
:RoomsDictUpdate: 更新聊天室字典
:SetRoom: 指定使用本聊天机器人的聊天室
"""
def __init__(self, username, password, domain):
# 设置队列，记录最近答复过的消息id，避免重复回答同一条消息，假定1分钟内@bot的人次不超过100
self.replied_queue = deque([], maxlen=100)
# 设置rocketchat的api参数
self.api = RocketChatAPI(settings={'username': username,
'password': password,
'domain': domain})
def _RoomsDict(self, _class='private'):
"""
获取聊天房间的字典，字典的key为聊天室名，value为聊天室id，聊天室不存在重名
:_class: 聊天室类型
- private: 私有聊天室
- public: 公共聊天室
"""
if _class == 'private':
rooms = self.api.get_private_rooms()
elif _class == 'public':
rooms = self.api.get_public_rooms()
else:
raise ValueError('请指定正确的_class')
fun = lambda d:tuple(d.values())
rooms = {fun(d)[0]:fun(d)[1] for d in rooms}
return rooms
def RoomsDictUpdate(self):
private_rooms = self._RoomsDict('private')
public_rooms = self._RoomsDict('public')
self.rooms_dict = dict(private_rooms, **public_rooms)
def SetRoom(self, room_name):
try:
self.room_id = self.rooms_dict[room_name]
except KeyError:
raise KeyError('房间名<%s>不存在' % room_name)
def History(self, oldest=None):
"""
获取当前聊天室的历史聊天记录
"""
history = self.api.get_private_room_history(self.room_id, oldest=oldest)
if history['success']:
history_msgs = history['messages']
return history_msgs
def NewMsgs(self):
"""
获取指定聊天室最近60秒的历史消息
"""
begin_CCT = datetime.now() + timedelta(seconds=-10)
begin_UTC = begin_CCT.astimezone(pytz.timezone('UTC'))
oldest = begin_UTC.strftime('%Y-%m-%dT%H:%M:%SZ')
msgs = self.History(oldest=oldest)
msgs.reverse() # 逆序排列
return msgs
def reply(self, ser):
if ser['_id'] not in self.replied_queue:
self.replied_queue.append(ser['_id'])
text_out = BotApi(ser['msg'])
text_out += ' @%s' % ser['name']
self.api.send_message(text_out, room_id)
def _Msgs2Bot(self, msgs: list):
"""
解析获取到的聊天记录，筛选出@xbot的记录，保留'_id'，'name'，'msg'这三个字段
"""
if len(msgs) > 0:
df_msgs = pd.DataFrame(msgs)
df_msgs['name'] = df_msgs['u'].map(lambda x:x['username'])
df_msgs_select = df_msgs[['_id', 'name', 'msg']]
self.msgs_at_bot = df_msgs_select[df_msgs_select['msg'].str.startswith('@xbot') |
df_msgs_select['msg'].str.endswith('@xbot')]
self.msgs_at_bot.apply(self.reply, axis=1)
def Msgs2Bot(self):
self._Msgs2Bot(self.NewMsgs())
def Run(self):
while True:
self.Msgs2Bot()
time.sleep(0.5)
if line.count('=') > 0:
chat_url = line.split('=')[1].strip().replace('"', '')
conf_file.close()
api = RocketChatAPI(settings={
'username': user,
'password': password,
'domain': chat_url
})
room_id = ''
for chat in chats.split(','):
for room in api.get_private_rooms():
if room['name'] == chat.strip():
room_id = room['id']
count = 0
for message in api.get_private_room_history(
room_id, oldest='2016-05-30T13:42:25.304Z',
count=2000)['messages']:
count = count + 1
message_ts = parser.parse(message['ts'])
utc = pytz.UTC
older_than = int((utc.localize(datetime.now()) -
message_ts).total_seconds() / 3600)
if older_than > 24:
# Удаляем сообщения старше 24-х часов
message_id = message['_id']
api.delete_private_room_message(room_id,
message_id=message_id)