misp-modules/misp_modules/helpers/cache.py

81 lines
2.3 KiB
Python
Raw Normal View History

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# MISP modules helper - cache
#
# Copyright (C) 2016 Alexandre Dulaunoy
# Copyright (C) 2016 CIRCL - Computer Incident Response Center Luxembourg
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
import redis
import hashlib
port = 6379
hostname = '127.0.0.1'
db = 5
def selftest(enable=True):
if not enable:
return False
r = redis.StrictRedis(host=hostname, port=port, db=db)
try:
2016-06-18 04:04:41 +02:00
r.ping()
except:
return 'Redis not running or not installed. Helper will be disabled.'
def get(modulename=None, query=None, value=None, debug=False):
2016-04-19 07:57:36 +02:00
if (modulename is None or query is None):
return False
r = redis.StrictRedis(host=hostname, port=port, db=db)
h = hashlib.sha1()
h.update(query.encode('UTF-8'))
hv = h.hexdigest()
2016-06-18 04:04:41 +02:00
key = "m:" + modulename + ":" + hv
if not r.exists(key):
if debug:
2016-06-18 04:04:41 +02:00
print("Key {} added in cache".format(key))
r.setex(key, 86400, value)
else:
if debug:
2016-06-18 04:04:41 +02:00
print("Cache hit with Key {}".format(key))
return r.get(key)
def flush():
r = redis.StrictRedis(host=hostname, port=port, db=db)
returncode = r.flushdb()
return returncode
if __name__ == "__main__":
import sys
if selftest() is not None:
sys.exit()
else:
2016-06-18 04:04:41 +02:00
print("Selftest ok")
v = get(modulename="testmodule", query="abcdef", value="barfoo", debug=True)
if v == b'barfoo':
2016-06-18 04:04:41 +02:00
print("Cache ok")
2016-04-19 07:57:36 +02:00
v = get(modulename="testmodule", query="abcdef")
2016-06-18 04:04:41 +02:00
print(v)
v = get(modulename="testmodule")
if (not v):
2016-06-18 04:04:41 +02:00
print("Failed ok")
if flush():
2016-06-18 04:04:41 +02:00
print("Cache flushed ok")