fixed output of coffee and primary banner

internal/sqlite/db.go

@@ -0,0 +1,143 @@
+package sqlite
+
+import (
+	"fmt"
+	"go.uber.org/zap"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
+	"gorm.io/gorm/logger"
+)
+
+var DB *gorm.DB
+
+// InitDB initializes the database connection
+func InitDB(dbPath string) (*gorm.DB, error) {
+	zap.L().Info("Initializing database", zap.String("path", dbPath))
+
+	// Check if the path is a file or directory
+	fileInfo, err := os.Stat(dbPath)
+	var finalDbPath string
+
+	// If path doesn't exist or is a directory
+	if os.IsNotExist(err) || (err == nil && fileInfo.IsDir()) {
+		// Treat as directory path
+		if err := os.MkdirAll(dbPath, 0755); err != nil {
+			zap.L().Error("Failed to create database directory", zap.Error(err))
+			return nil, fmt.Errorf("failed to create database directory: %w", err)
+		}
+		finalDbPath = filepath.Join(dbPath, "dehashed.sqlite")
+	} else {
+		// Treat as file path
+		// Ensure the directory exists
+		dir := filepath.Dir(dbPath)
+		if err := os.MkdirAll(dir, 0755); err != nil {
+			zap.L().Error("Failed to create parent directory for database", zap.Error(err))
+			return nil, fmt.Errorf("failed to create parent directory for database: %w", err)
+		}
+		finalDbPath = dbPath
+	}
+
+	zap.L().Info("Opening database", zap.String("finalPath", finalDbPath))
+	db, err := gorm.Open(sqlite.Open(finalDbPath), &gorm.Config{
+		Logger: logger.Default.LogMode(logger.Silent),
+	})
+	if err != nil {
+		zap.L().Error("Failed to connect to database", zap.Error(err))
+		return nil, fmt.Errorf("failed to connect to database: %w", err)
+	}
+
+	// Auto migrate your models
+	err = db.AutoMigrate(&Result{}, &Creds{}, &QueryOptions{}, &Creds{}, &WhoisRecord{}, &SubdomainRecord{},
+		&HistoryRecord{}, &LookupResult{}, &HunterDomainData{}, &HunterEmail{}, &PersonData{})
+	if err != nil {
+		zap.L().Error("Failed to migrate database", zap.Error(err))
+		return nil, fmt.Errorf("failed to migrate database: %w", err)
+	}
+
+	DB = db
+	return db, nil
+}
+
+// GetDB returns the database connection
+func GetDB() *gorm.DB {
+	if DB == nil {
+		zap.L().Error("database not initialized")
+		fmt.Println("sqlite database not initialized")
+		os.Exit(1)
+	}
+	return DB
+}
+
+type Table int64
+
+const (
+	ResultsTable Table = iota
+	RunsTable
+	CredsTable
+	WhoIsTable
+	SubdomainsTable
+	HistoryTable
+	LookupTable
+	HunterDomainTable
+	HunterEmailTable
+	PersonTable
+	UnknownTable
+)
+
+func GetTable(userInput string) Table {
+	switch strings.ToLower(userInput) {
+	case "results":
+		return ResultsTable
+	case "runs":
+		return RunsTable
+	case "creds":
+		return CredsTable
+	case "whois":
+		return WhoIsTable
+	case "subdomains":
+		return SubdomainsTable
+	case "history":
+		return HistoryTable
+	case "lookup":
+		return LookupTable
+	case "hunter_domain":
+		return HunterDomainTable
+	case "hunter_email":
+		return HunterEmailTable
+	case "person":
+		return PersonTable
+	default:
+		return UnknownTable
+	}
+}
+
+func (t Table) Object() interface{} {
+	switch t {
+	case ResultsTable:
+		return Result{}
+	case RunsTable:
+		return QueryOptions{}
+	case CredsTable:
+		return Creds{}
+	case WhoIsTable:
+		return WhoisRecord{}
+	case SubdomainsTable:
+		return SubdomainRecord{}
+	case HistoryTable:
+		return HistoryRecord{}
+	case LookupTable:
+		return LookupResult{}
+	case HunterDomainTable:
+		return HunterDomainData{}
+	case HunterEmailTable:
+		return HunterEmail{}
+	case PersonTable:
+		return PersonData{}
+	default:
+		return nil
+	}
+}