Skip to content

Commit 32d1383

Browse files
rohitcoderrohitcoder
committed
Added support for PostgreSQL
1 parent 0c78605 commit 32d1383

File tree

5 files changed

+124
-25
lines changed

5 files changed

+124
-25
lines changed

‎connection.yml.sample‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
notify:
22
slack:
33
webhook_url: https://hooks.slack.com/services/T0XXXXXXXXXXX/BXXXXXXXX/1CIyXXXXXXXXXXXXXXX
4-
4+
55
sources:
66
redis:
77
redis_example:
@@ -35,6 +35,13 @@ sources:
3535
user: YOUR_MYSQL_USERNAME
3636
password: YOUR_MYSQL_PASSWORD
3737
database: YOUR_MYSQL_DATABASE_NAME
38+
postgresql: # New PostgreSQL configuration
39+
postgresql_profile:
40+
host: YOUR_POSTGRESQL_HOST
41+
port: YOUR_POSTGRESQL_PORT
42+
user: YOUR_POSTGRESQL_USERNAME
43+
password: YOUR_POSTGRESQL_PASSWORD
44+
database: YOUR_POSTGRESQL_DATABASE_NAME
3845
fs:
3946
fs_example:
4047
path: /path/to/your/filesystem/directory

‎hawk_scanner/commands/mysql.py‎

Lines changed: 20 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1,31 +1,29 @@
1-
import mysql.connector
1+
import psycopg2
22
from hawk_scanner.internals import system
33
import re
44
from rich.console import Console
55
from rich.table import Table
66

77
console = Console()
88

9-
def connect_mysql(host, port, user, password, database):
9+
def connect_postgresql(host, port, user, password, database):
1010
try:
11-
conn = mysql.connector.connect(
11+
conn = psycopg2.connect(
1212
host=host,
1313
port=port,
1414
user=user,
1515
password=password,
1616
database=database
1717
)
18-
if conn.is_connected():
19-
system.print_info(f"Connected to MySQL database at {host}")
18+
if conn:
19+
system.print_info(f"Connected to PostgreSQL database at {host}")
2020
return conn
21-
else:
22-
system.print_error(f"Failed to connect to MySQL database at {host}")
2321
except Exception as e:
24-
system.print_error(f"Failed to connect to MySQL database at {host} with error: {e}")
22+
system.print_error(f"Failed to connect to PostgreSQL database at {host} with error: {e}")
2523

2624
def check_data_patterns(conn, patterns, profile_name, database_name):
2725
cursor = conn.cursor()
28-
cursor.execute("SHOW TABLES")
26+
cursor.execute("SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'")
2927
tables = [table[0] for table in cursor.fetchall()]
3028

3129
table_count = 1
@@ -44,15 +42,15 @@ def check_data_patterns(conn, patterns, profile_name, database_name):
4442
if matches:
4543
for match in matches:
4644
results.append({
47-
'host': conn._host,
45+
'host': conn.dsn,
4846
'database': database_name,
4947
'table': table,
5048
'column': column,
5149
'pattern_name': match['pattern_name'],
5250
'matches': match['matches'],
5351
'sample_text': match['sample_text'],
5452
'profile': profile_name,
55-
'data_source': 'mysql'
53+
'data_source': 'postgresql'
5654
})
5755

5856
data_count += 1
@@ -64,33 +62,33 @@ def check_data_patterns(conn, patterns, profile_name, database_name):
6462

6563
def execute(args):
6664
results = []
67-
system.print_info(f"Running Checks for MySQL Sources")
65+
system.print_info(f"Running Checks for PostgreSQL Sources")
6866
connections = system.get_connection()
6967

7068
if 'sources' in connections:
7169
sources_config = connections['sources']
72-
mysql_config = sources_config.get('mysql')
70+
postgresql_config = sources_config.get('postgresql')
7371

74-
if mysql_config:
72+
if postgresql_config:
7573
patterns = system.get_fingerprint_file()
7674

77-
for key, config in mysql_config.items():
75+
for key, config in postgresql_config.items():
7876
host = config.get('host')
7977
user = config.get('user')
80-
port = config.get('port', 3306) # default port
78+
port = config.get('port', 5432) # default port for PostgreSQL
8179
password = config.get('password')
8280
database = config.get('database')
8381

8482
if host and user and password and database:
85-
system.print_info(f"Checking MySQL Profile {key} and database {database}")
86-
conn = connect_mysql(host, port, user, password, database)
83+
system.print_info(f"Checking PostgreSQL Profile {key} and database {database}")
84+
conn = connect_postgresql(host, port, user, password, database)
8785
if conn:
88-
results = check_data_patterns(conn, patterns, key, database)
86+
results += check_data_patterns(conn, patterns, key, database)
8987
conn.close()
9088
else:
91-
system.print_error(f"Incomplete MySQL configuration for key: {key}")
89+
system.print_error(f"Incomplete PostgreSQL configuration for key: {key}")
9290
else:
93-
system.print_error("No MySQL connection details found in connection.yml")
91+
system.print_error("No PostgreSQL connection details found in connection.yml")
9492
else:
9593
system.print_error("No 'sources' section found in connection.yml")
96-
return results
94+
return results

‎hawk_scanner/commands/postgresql.py‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
import psycopg2
2+
from hawk_scanner.internals import system
3+
from rich.console import Console
4+
from rich.table import Table
5+
6+
console = Console()
7+
8+
def connect_postgresql(host, port, user, password, database):
9+
try:
10+
conn = psycopg2.connect(
11+
host=host,
12+
port=port,
13+
user=user,
14+
password=password,
15+
database=database
16+
)
17+
if conn:
18+
system.print_info(f"Connected to PostgreSQL database at {host}")
19+
return conn
20+
except Exception as e:
21+
system.print_error(f"Failed to connect to PostgreSQL database at {host} with error: {e}")
22+
23+
def check_data_patterns(conn, patterns, profile_name, database_name):
24+
cursor = conn.cursor()
25+
cursor.execute("SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'")
26+
tables = [table[0] for table in cursor.fetchall()]
27+
28+
table_count = 1
29+
30+
results = []
31+
for table in tables:
32+
cursor.execute(f"SELECT * FROM {table}")
33+
columns = [column[0] for column in cursor.description]
34+
35+
data_count = 1
36+
for row in cursor.fetchall():
37+
for column, value in zip(columns, row):
38+
if value:
39+
value_str = str(value)
40+
matches = system.match_strings(value_str)
41+
if matches:
42+
for match in matches:
43+
results.append({
44+
'host': conn.dsn,
45+
'database': database_name,
46+
'table': table,
47+
'column': column,
48+
'pattern_name': match['pattern_name'],
49+
'matches': match['matches'],
50+
'sample_text': match['sample_text'],
51+
'profile': profile_name,
52+
'data_source': 'postgresql'
53+
})
54+
55+
data_count += 1
56+
57+
table_count += 1
58+
59+
cursor.close()
60+
return results
61+
62+
def execute(args):
63+
results = []
64+
system.print_info(f"Running Checks for PostgreSQL Sources")
65+
connections = system.get_connection()
66+
67+
if 'sources' in connections:
68+
sources_config = connections['sources']
69+
postgresql_config = sources_config.get('postgresql')
70+
71+
if postgresql_config:
72+
patterns = system.get_fingerprint_file()
73+
74+
for key, config in postgresql_config.items():
75+
host = config.get('host')
76+
user = config.get('user')
77+
port = config.get('port', 5432) # default port for PostgreSQL
78+
password = config.get('password')
79+
database = config.get('database')
80+
81+
if host and user and password and database:
82+
system.print_info(f"Checking PostgreSQL Profile {key} and database {database}")
83+
conn = connect_postgresql(host, port, user, password, database)
84+
if conn:
85+
results += check_data_patterns(conn, patterns, key, database)
86+
conn.close()
87+
else:
88+
system.print_error(f"Incomplete PostgreSQL configuration for key: {key}")
89+
else:
90+
system.print_error("No PostgreSQL connection details found in connection.yml")
91+
else:
92+
system.print_error("No 'sources' section found in connection.yml")
93+
return results

‎hawk_scanner/main.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ def clear_screen():
2121
console = Console()
2222

2323
## Now separate the results by data_source
24-
data_sources = ['s3', 'mysql', 'redis', 'firebase', 'gcs', 'fs']
24+
data_sources = ['s3', 'mysql', 'redis', 'firebase', 'gcs', 'fs', 'postgresql']
2525

2626
def load_command_module(command):
2727
try:

‎requirements.txt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,4 +5,5 @@ mysql-connector-python
55
redis
66
firebase-admin
77
google-cloud-core
8-
google-cloud-storage
8+
google-cloud-storage
9+
psycopg2

0 commit comments

Comments
 (0)