From 57cb6573108019e21fe171385675df22d6fe3fb6 Mon Sep 17 00:00:00 2001 From: Bowen Chen Date: Sun, 21 Jul 2019 04:50:03 -0700 Subject: [PATCH] upgrade pipot from python2 to python3 (#30) --- config_parser.py | 2 +- database.py | 91 ---------------------------------- mod_config/controllers.py | 9 ++-- mod_config/models.py | 29 +++++------ mod_honeypot/controllers.py | 2 +- mod_honeypot/models.py | 24 ++++----- tests/testServiceManagement.py | 3 +- 7 files changed, 36 insertions(+), 124 deletions(-) diff --git a/config_parser.py b/config_parser.py index 24df2f3..b1e2097 100644 --- a/config_parser.py +++ b/config_parser.py @@ -12,7 +12,7 @@ def parse_config(obj): :rtype: dict """ config = {} - if isinstance(obj, (str, unicode)): + if isinstance(obj, (bytes, str)): obj = import_string(obj) for key in dir(obj): if key.isupper(): diff --git a/database.py b/database.py index 8c52f26..1e531bd 100644 --- a/database.py +++ b/database.py @@ -4,7 +4,6 @@ from sqlalchemy import create_engine from sqlalchemy.ext.declarative import declarative_base, DeclarativeMeta from sqlalchemy.orm import scoped_session, sessionmaker -from sqlalchemy.sql.sqltypes import SchemaType, Enum, TypeDecorator class DeclarativeABCMeta(DeclarativeMeta, ABCMeta): @@ -42,93 +41,3 @@ def create_session(db_string, drop_tables=False): Base.metadata.create_all(bind=db_engine) return db_session - - -class EnumSymbol(object): - """Define a fixed symbol tied to a parent class.""" - - def __init__(self, cls_, name, value, description): - self.cls_ = cls_ - self.name = name - self.value = value - self.description = description - - def __reduce__(self): - """Allow unpickling to return the symbol linked to the DeclEnum - class.""" - return getattr, (self.cls_, self.name) - - def __iter__(self): - return iter([self.value, self.description]) - - def __repr__(self): - return "<%s>" % self.name - - -class EnumMeta(type): - """Generate new DeclEnum classes.""" - - def __init__(cls, classname, bases, dict_): - cls._reg = reg = cls._reg.copy() - for k, v in dict_.items(): - if isinstance(v, tuple): - sym = reg[v[0]] = EnumSymbol(cls, k, *v) - setattr(cls, k, sym) - return type.__init__(cls, classname, bases, dict_) - - def __iter__(cls): - return iter(cls._reg.values()) - - -class DeclEnum(object): - """Declarative enumeration.""" - - __metaclass__ = EnumMeta - - _reg = {} - - @classmethod - def from_string(cls, value): - try: - return cls._reg[value] - except KeyError: - raise ValueError( - "Invalid value for %r: %r" % - (cls.__name__, value) - ) - - @classmethod - def values(cls): - return cls._reg.keys() - - @classmethod - def db_type(cls): - return DeclEnumType(cls) - - -class DeclEnumType(SchemaType, TypeDecorator): - def __init__(self, enum): - self.enum = enum - self.impl = Enum( - *enum.values(), - name="ck%s" % re.sub( - '([A-Z])', - lambda m: "_" + m.group(1).lower(), - enum.__name__) - ) - - def _set_table(self, table, column): - self.impl._set_table(table, column) - - def copy(self): - return DeclEnumType(self.enum) - - def process_bind_param(self, value, dialect): - if value is None: - return None - return value.value - - def process_result_value(self, value, dialect): - if value is None: - return None - return self.enum.from_string(value.strip()) diff --git a/mod_config/controllers.py b/mod_config/controllers.py index a73e16a..8acd5a8 100644 --- a/mod_config/controllers.py +++ b/mod_config/controllers.py @@ -235,8 +235,8 @@ def notifications_ajax(action): @template_renderer() def data_processing(): form = RuleForm(request.form) - form.action.choices = [(key, value) for key, value in Actions] - form.condition.choices = [(key, value) for key, value in Conditions] + form.action.choices = [(item.name, item.value) for item in Actions] + form.condition.choices = [(item.name, item.value) for item in Conditions] notification_services = [(n.id, n.name) for n in Notification.query.all()] form.notification_id.choices = notification_services form.service_id.choices = [(s.id, s.name) for s in Service.query.all()] @@ -245,8 +245,9 @@ def data_processing(): rule = Rule( form.service_id.data, form.notification_id.data, form.notification_config.data, - Conditions.from_string(form.condition.data), form.level.data, - Actions.from_string(form.action.data) + Conditions[form.condition.data], + form.level.data, + Actions[form.action.data] ) g.db.add(rule) g.db.commit() diff --git a/mod_config/models.py b/mod_config/models.py index 422d1a0..acadbaa 100644 --- a/mod_config/models.py +++ b/mod_config/models.py @@ -1,11 +1,12 @@ import json import os +import enum from sqlalchemy import Column, Integer, String, Text, ForeignKey, Boolean, \ - UniqueConstraint + UniqueConstraint, Enum from sqlalchemy.orm import relationship -from database import Base, DeclEnum +from database import Base class Service(Base): @@ -52,18 +53,18 @@ def get_file(self, temp_folder=False): ) -class Actions(DeclEnum): - drop = "drop", "Drop" - store = "store", "Store" +class Actions(enum.Enum): + drop = "drop" + store = "store" -class Conditions(DeclEnum): - st = '<', '<' - gt = '>', '>' - eq = '==', '==' - se = '<=', '<=' - ge = '>=', '>=' - ne = '!=', '!=' +class Conditions(enum.Enum): + st = '<' + gt = '>' + eq = '==' + se = '<=' + ge = '>=' + ne = '!=' class Rule(Base): @@ -80,9 +81,9 @@ class Rule(Base): ondelete="CASCADE"), nullable=True) notification_config = Column(Text) - condition = Column(Conditions.db_type()) + condition = Column(Enum(Conditions)) level = Column(Integer) - action = Column(Actions.db_type()) + action = Column(Enum(Actions)) UniqueConstraint('service_id', 'notification_id', 'condition') service = relationship(Service) notification = relationship(Notification) diff --git a/mod_honeypot/controllers.py b/mod_honeypot/controllers.py index 94a913c..ee9b587 100644 --- a/mod_honeypot/controllers.py +++ b/mod_honeypot/controllers.py @@ -18,7 +18,7 @@ from mod_config.models import Service from mod_honeypot.forms import NewDeploymentForm, ModifyProfileForm, \ NewProfileForm, ServiceProfileForm -from models import Profile, ProfileService, Deployment, PiPotReport, \ +from mod_honeypot.models import Profile, ProfileService, Deployment, PiPotReport, \ PiModels, CollectorTypes mod_honeypot = Blueprint('honeypot', __name__) diff --git a/mod_honeypot/models.py b/mod_honeypot/models.py index ca6cbf3..0f0c2ad 100644 --- a/mod_honeypot/models.py +++ b/mod_honeypot/models.py @@ -1,10 +1,11 @@ import datetime import json import os +import enum -from sqlalchemy import Column, Integer, String, Text, ForeignKey, orm, Boolean +from sqlalchemy import Column, Integer, String, Text, ForeignKey, orm, Boolean, Enum -from database import Base, DeclEnum +from database import Base from pipot.services.IService import IModel @@ -62,16 +63,15 @@ def get_service_config(self): return json.loads(self.service_configuration) -class PiModels(DeclEnum): - one = "one", "Raspberry Pi 1 model B" - two = "two", "Raspberry Pi 2" - three = "three", "Raspberry Pi 3" +class PiModels(enum.Enum): + one = "one" + two = "two" + three = "three" -class CollectorTypes(DeclEnum): - udp = "udp", "UDP" - tcp = "tcp", "TCP (SSL)" - +class CollectorTypes(enum.Enum): + udp = "udp" + tcp = "tcp" class Deployment(Base): __tablename__ = 'deployment' @@ -86,14 +86,14 @@ class Deployment(Base): instance_key = Column(String(20)) mac_key = Column(String(32)) encryption_key = Column(String(32)) - rpi_model = Column(PiModels.db_type()) + rpi_model = Column(Enum(PiModels)) server_ip = Column(String(46)) # IPv6 proof interface = Column(String(20)) wlan_config = Column(Text()) hostname = Column(String(64)) rootpw = Column(String(50)) debug = Column(Boolean()) - collector_type = Column(CollectorTypes.db_type()) + collector_type = Column(Enum(CollectorTypes)) def __init__(self, name, profile_id, instance_key, mac_key, encryption_key, rpi_model, server_ip, interface, diff --git a/tests/testServiceManagement.py b/tests/testServiceManagement.py index 9808c58..94416e8 100644 --- a/tests/testServiceManagement.py +++ b/tests/testServiceManagement.py @@ -3,6 +3,7 @@ import mock import unittest import json +import codecs from mock import patch from functools import wraps from werkzeug.datastructures import FileStorage @@ -30,7 +31,7 @@ def add_and_remove_service(self, service_name, service_file_name): # upload the service file # service_name = 'TelnetService' # service_file_name = service_name + '.py' - service_file = open(os.path.join(test_dir, 'testFiles', service_file_name), 'r') + service_file = codecs.open(os.path.join(test_dir, 'testFiles', service_file_name), 'rb') # service_file = FileStorage(service_file) with self.app.test_client() as client: data = dict(