93 lines
3.4 KiB
Python
93 lines
3.4 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Migration: add security columns.
|
|
users — totp_secret, totp_enabled
|
|
schwab_connections — refresh_token_expires_at
|
|
teller_enrollments — widen access_token VARCHAR(128) → TEXT (needed for Fernet-encrypted values)
|
|
schwab_connections — widen access_token / refresh_token TEXT (already TEXT; no-op guard included)
|
|
|
|
Run once: python scripts/add_security_columns.py
|
|
Safe to re-run — skips columns that already exist or are already wide enough.
|
|
"""
|
|
import sys, os
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from app import create_app
|
|
from app.extensions import db
|
|
|
|
app = create_app()
|
|
|
|
COLUMNS = [
|
|
("users", "totp_secret",
|
|
"VARCHAR(64) NULL DEFAULT NULL"),
|
|
("users", "totp_enabled",
|
|
"TINYINT(1) NOT NULL DEFAULT 0"),
|
|
("schwab_connections", "refresh_token_expires_at",
|
|
"DATETIME NULL DEFAULT NULL"),
|
|
]
|
|
|
|
CREATE_TABLES = [
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS audit_logs (
|
|
id INT AUTO_INCREMENT PRIMARY KEY,
|
|
timestamp DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
|
action VARCHAR(64) NOT NULL,
|
|
description VARCHAR(255),
|
|
ip_address VARCHAR(45),
|
|
INDEX ix_audit_logs_timestamp (timestamp),
|
|
INDEX ix_audit_logs_action (action)
|
|
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
|
|
"""
|
|
]
|
|
|
|
# Columns whose type must be widened to TEXT to hold Fernet-encrypted values.
|
|
# ALTER COLUMN is safe to re-run — widening TEXT→TEXT is a no-op in MySQL.
|
|
WIDEN_COLUMNS = [
|
|
("teller_enrollments", "access_token", "TEXT NOT NULL"),
|
|
("schwab_connections", "access_token", "TEXT NOT NULL"),
|
|
("schwab_connections", "refresh_token", "TEXT NOT NULL"),
|
|
]
|
|
|
|
with app.app_context():
|
|
with db.engine.connect() as conn:
|
|
for table, column, definition in WIDEN_COLUMNS:
|
|
current_type = conn.execute(db.text(
|
|
"SELECT DATA_TYPE FROM information_schema.columns "
|
|
"WHERE table_schema = DATABASE() "
|
|
"AND table_name = :t AND column_name = :c"
|
|
), {"t": table, "c": column}).scalar()
|
|
|
|
if current_type and current_type.lower() == 'text':
|
|
print(f" {table}.{column} already TEXT — skipped.")
|
|
elif current_type:
|
|
conn.execute(db.text(
|
|
f"ALTER TABLE `{table}` MODIFY COLUMN `{column}` {definition}"
|
|
))
|
|
conn.commit()
|
|
print(f" Widened {table}.{column} to TEXT.")
|
|
else:
|
|
print(f" {table}.{column} not found — skipped.")
|
|
|
|
for table, column, definition in COLUMNS:
|
|
exists = conn.execute(db.text(
|
|
"SELECT COUNT(*) FROM information_schema.columns "
|
|
"WHERE table_schema = DATABASE() "
|
|
"AND table_name = :t AND column_name = :c"
|
|
), {"t": table, "c": column}).scalar()
|
|
|
|
if exists:
|
|
print(f" {table}.{column} already exists — skipped.")
|
|
else:
|
|
conn.execute(db.text(
|
|
f"ALTER TABLE {table} ADD COLUMN {column} {definition}"
|
|
))
|
|
conn.commit()
|
|
print(f" Added {table}.{column}.")
|
|
|
|
for sql in CREATE_TABLES:
|
|
conn.execute(db.text(sql))
|
|
conn.commit()
|
|
print(" Created audit_logs table (or already exists).")
|
|
|
|
print("Done.")
|