135 lines
5.0 KiB
Python
135 lines
5.0 KiB
Python
import sqlite3
|
|
from datetime import datetime
|
|
import os
|
|
|
|
DB_FILE = "sms_app.db"
|
|
|
|
class DBHandler:
|
|
def __init__(self, db_path=DB_FILE):
|
|
self.db_path = db_path
|
|
self._initialize_db()
|
|
|
|
def _get_connection(self):
|
|
return sqlite3.connect(self.db_path)
|
|
|
|
def _initialize_db(self):
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
# Create contacts table
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS contacts (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
name TEXT NOT NULL,
|
|
phone TEXT UNIQUE NOT NULL,
|
|
public_key TEXT
|
|
)
|
|
''')
|
|
# Create messages table
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS messages (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
phone TEXT NOT NULL,
|
|
text TEXT NOT NULL,
|
|
date TEXT NOT NULL,
|
|
is_secure INTEGER NOT NULL,
|
|
status TEXT NOT NULL,
|
|
FOREIGN KEY(phone) REFERENCES contacts(phone)
|
|
)
|
|
''')
|
|
# Create settings table for our own keypair if needed (though usually we save keys to files)
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS settings (
|
|
key TEXT PRIMARY KEY,
|
|
value TEXT
|
|
)
|
|
''')
|
|
conn.commit()
|
|
|
|
# --- Contacts Methods ---
|
|
def add_contact(self, name, phone, public_key=None):
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
try:
|
|
cursor.execute('INSERT INTO contacts (name, phone, public_key) VALUES (?, ?, ?)',
|
|
(name, phone, public_key))
|
|
conn.commit()
|
|
return True
|
|
except sqlite3.IntegrityError:
|
|
return False # Phone number already exists
|
|
|
|
def update_contact_key(self, phone, public_key):
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('UPDATE contacts SET public_key = ? WHERE phone = ?', (public_key, phone))
|
|
conn.commit()
|
|
|
|
def get_contact(self, phone):
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('SELECT name, phone, public_key FROM contacts WHERE phone = ?', (phone,))
|
|
return cursor.fetchone()
|
|
|
|
def get_all_contacts(self):
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('SELECT name, phone, public_key FROM contacts')
|
|
return cursor.fetchall()
|
|
|
|
def set_contact_name_if_not_exists(self, phone, name):
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('SELECT phone FROM contacts WHERE phone = ?', (phone,))
|
|
if not cursor.fetchone():
|
|
cursor.execute('INSERT INTO contacts (name, phone, public_key) VALUES (?, ?, ?)', (name, phone, None))
|
|
conn.commit()
|
|
|
|
# --- Messages Methods ---
|
|
def add_message(self, phone, text, is_secure, status):
|
|
date_str = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('''
|
|
INSERT INTO messages (phone, text, date, is_secure, status)
|
|
VALUES (?, ?, ?, ?, ?)
|
|
''', (phone, text, date_str, int(is_secure), status))
|
|
conn.commit()
|
|
return cursor.lastrowid
|
|
|
|
def get_messages_for_contact(self, phone):
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('''
|
|
SELECT id, phone, text, date, is_secure, status
|
|
FROM messages
|
|
WHERE phone = ?
|
|
ORDER BY date ASC
|
|
''', (phone,))
|
|
return cursor.fetchall()
|
|
|
|
def delete_message(self, msg_id):
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('DELETE FROM messages WHERE id = ?', (msg_id,))
|
|
conn.commit()
|
|
|
|
# --- Settings Methods ---
|
|
def get_setting(self, key):
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('SELECT value FROM settings WHERE key = ?', (key,))
|
|
row = cursor.fetchone()
|
|
return row[0] if row else None
|
|
|
|
def set_setting(self, key, value):
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)', (key, value))
|
|
conn.commit()
|
|
|
|
if __name__ == "__main__":
|
|
db = DBHandler("test_sms.db")
|
|
db.add_contact("Alice", "+1234567890", "PUBLIC_KEY_CONTENT")
|
|
db.add_message("+1234567890", "Hello Alice!", is_secure=False, status="sent")
|
|
print("Database OK.")
|
|
os.remove("test_sms.db")
|