[#25][FIX] Logic error in column verification

dfir-iris · Jan 2, 2022 · 9f27ecb · 9f27ecb
1 parent fdb2a16
commit 9f27ecb
Showing 1 changed file with 41 additions and 41 deletions.
diff --git a/source/app/alembic/versions/6a3b3b627d45_add_ioc_type.py b/source/app/alembic/versions/6a3b3b627d45_add_ioc_type.py
@@ -47,43 +47,43 @@ def upgrade():
             local_cols=["ioc_type_id"],
             remote_cols=["type_id"])
 
-        if _table_has_column('ioc', 'ioc_type'):
-            # Set schema and make migration of data
-            t_ioc = sa.Table(
-                'ioc',
-                sa.MetaData(),
-                sa.Column('ioc_id', sa.Integer, primary_key=True),
-                sa.Column('ioc_value', sa.Text),  # Old column.
-                sa.Column('ioc_type', sa.Unicode(length=50)),  # Two new columns.
-                sa.Column('ioc_type_id', sa.ForeignKey('ioc_type.type_id')),
-                sa.Column('ioc_tags', sa.Text),
-                sa.Column('user_id', sa.ForeignKey('user.id')),
-                sa.Column('ioc_misp', sa.Text),
-                sa.Column('ioc_tlp_id', sa.ForeignKey('tlp.tlp_id'))
+    if _table_has_column('ioc', 'ioc_type'):
+        # Set schema and make migration of data
+        t_ioc = sa.Table(
+            'ioc',
+            sa.MetaData(),
+            sa.Column('ioc_id', sa.Integer, primary_key=True),
+            sa.Column('ioc_value', sa.Text),
+            sa.Column('ioc_type', sa.Unicode(length=50)),
+            sa.Column('ioc_type_id', sa.ForeignKey('ioc_type.type_id')),
+            sa.Column('ioc_tags', sa.Text),
+            sa.Column('user_id', sa.ForeignKey('user.id')),
+            sa.Column('ioc_misp', sa.Text),
+            sa.Column('ioc_tlp_id', sa.ForeignKey('tlp.tlp_id'))
+        )
+        to_update = [('Domain', 'domain'), ('IP', 'ip-any'), ('Hash', 'other'), ('File', 'filename'),
+                     ('Path', 'file-path'), ('Account', 'account'), ("Other", 'other')]
+
+        # Migrate existing IOCs
+        for src_up, dst_up in to_update:
+            conn = op.get_bind()
+            res = conn.execute(f"select ioc_id from ioc where ioc_type = '{src_up}';")
+            results = res.fetchall()
+            res = conn.execute(f"select type_id from ioc_type where type_name = '{dst_up}';")
+            e_info = res.fetchall()
+
+            if e_info:
+                domain_id = e_info[0][0]
+
+                for res in results:
+                    conn.execute(t_ioc.update().where(t_ioc.c.ioc_id == res[0]).values(
+                        ioc_type_id=domain_id
+                    ))
+
+            op.drop_column(
+                table_name='ioc',
+                column_name='ioc_type'
             )
-            to_update = [('Domain', 'domain'), ('IP', 'ip-any'), ('Hash', 'other'), ('File', 'filename'),
-                         ('Path', 'file-path'), ('Account', 'account'), ("Other", 'other')]
-
-            # Migrate existing IOCs
-            for src_up, dst_up in to_update:
-                conn = op.get_bind()
-                res = conn.execute(f"select ioc_id from ioc where ioc_type = '{src_up}';")
-                results = res.fetchall()
-                res = conn.execute(f"select type_id from ioc_type where type_name = '{dst_up}';")
-                e_info = res.fetchall()
-
-                if e_info:
-                    domain_id = e_info[0][0]
-
-                    for res in results:
-                        conn.execute(t_ioc.update().where(t_ioc.c.ioc_id == res[0]).values(
-                            ioc_type_id=domain_id
-                        ))
-
-                op.drop_column(
-                    table_name='ioc',
-                    column_name='ioc_type'
-                )
 
     pass
 
@@ -97,9 +97,9 @@ def _table_has_column(table, column):
     engine = engine_from_config(
         config.get_section(config.config_ini_section), prefix='sqlalchemy.')
     insp = reflection.Inspector.from_engine(engine)
-    has_column = False
+
     for col in insp.get_columns(table):
-        if column not in col['name']:
-            continue
-        has_column = True
-    return has_column
+        if column == col['name']:
+            return True
+
+    return False