|
20 | 20 | @hug.get('/test') |
21 | 21 | def test_connect(): |
22 | 22 | """Test connection to db.""" |
23 | | - psycopg2.connect("dbname='t' user='t' host='localhost' password='test'") |
| 23 | + # Get database credentials from configs.py or environment variables |
| 24 | + try: |
| 25 | + from configs import db_name, db_user, db_host, db_password |
| 26 | + except ImportError: |
| 27 | + import os |
| 28 | + db_name = os.getenv('DB_NAME', 't') |
| 29 | + db_user = os.getenv('DB_USER', 't') |
| 30 | + db_host = os.getenv('DB_HOST', 'localhost') |
| 31 | + db_password = os.getenv('DB_PASSWORD', '') |
| 32 | + if not db_password: |
| 33 | + return {'error': 'Database password not configured. Set DB_PASSWORD environment variable or configs.py'} |
| 34 | + |
| 35 | + psycopg2.connect(f"dbname='{db_name}' user='{db_user}' host='{db_host}' password='{db_password}'") |
24 | 36 | return ('connected successfully to db! ready for queries.') |
25 | 37 |
|
26 | 38 |
|
27 | 39 | @hug.get('/checktable') |
28 | 40 | def test_write(user='t', table='testtable'): |
29 | 41 | """Test write to DB.""" |
30 | | - conn = psycopg2.connect("dbname='t' user='t' host='localhost' password='test'") |
| 42 | + # Get database credentials from configs.py or environment variables |
| 43 | + try: |
| 44 | + from configs import db_name, db_user, db_host, db_password |
| 45 | + except ImportError: |
| 46 | + import os |
| 47 | + db_name = os.getenv('DB_NAME', 't') |
| 48 | + db_user = os.getenv('DB_USER', 't') |
| 49 | + db_host = os.getenv('DB_HOST', 'localhost') |
| 50 | + db_password = os.getenv('DB_PASSWORD', '') |
| 51 | + if not db_password: |
| 52 | + return {'error': 'Database password not configured. Set DB_PASSWORD environment variable or configs.py'} |
| 53 | + |
| 54 | + conn = psycopg2.connect(f"dbname='{db_name}' user='{db_user}' host='{db_host}' password='{db_password}'") |
31 | 55 | print('connected successfully to db! ready for queries.') |
32 | 56 | cur = conn.cursor() |
33 | | - cur.execute("select exists(select relname from pg_class where relname='" + table + "')") |
| 57 | + # Fix SQL injection vulnerability by using parameterized queries |
| 58 | + cur.execute("SELECT exists(SELECT relname FROM pg_class WHERE relname=%s)", (table,)) |
34 | 59 | exists = cur.fetchone()[0] |
35 | 60 | print(exists) |
36 | 61 | cur.close() |
| 62 | + conn.close() |
37 | 63 | if exists: |
38 | 64 | return 'THIS TABLE EXISTS' |
39 | 65 | else: |
|
0 commit comments