Adaptix-Framework-AdaptixC2/AdaptixServer/core/database/database.go

package database

import (
	"database/sql"

	_ "github.com/mattn/go-sqlite3"
)

type DBMS struct {
	database *sql.DB
	exists   bool
}

func NewDatabase(dbPath string) (*DBMS, error) {
	var err error

	dbms := &DBMS{
		exists: true,
	}

	// Enable WAL mode and other performance optimizations via connection string
	connStr := dbPath + "?_journal_mode=WAL&_synchronous=NORMAL&_busy_timeout=10000&_cache_size=-64000"
	dbms.database, err = sql.Open("sqlite3", connStr)
	if err != nil {
		dbms.exists = false
	}

	if dbms.exists {
		dbms.database.SetMaxOpenConns(1)
		dbms.database.SetMaxIdleConns(1)

		err = dbms.DatabaseInit()
		if err != nil {
			dbms.exists = false
		}
	}
	return dbms, err
}

func (dbms *DBMS) DatabaseInit() error {
	var (
		err              error
		createTableQuery string
	)

	createTableQuery = `CREATE TABLE IF NOT EXISTS "Listeners" (
    	"ListenerName" TEXT NOT NULL UNIQUE, 
    	"ListenerRegName" TEXT NOT NULL,
    	"ListenerConfig" TEXT NOT NULL,
    	"CreateTime" BIGINT,
    	"Watermark" TEXT NOT NULL,
    	"CustomData" BLOB
    );`
	_, err = dbms.database.Exec(createTableQuery)

	createTableQuery = `CREATE TABLE IF NOT EXISTS "Chat" (
    	"Id" INTEGER PRIMARY KEY AUTOINCREMENT,
    	"Username" TEXT NOT NULL,
    	"Message" TEXT NOT NULL,
    	"Date" BIGINT
    );`
	_, err = dbms.database.Exec(createTableQuery)

	createTableQuery = `CREATE TABLE IF NOT EXISTS "Downloads" (
    	"FileId" TEXT NOT NULL UNIQUE, 
    	"AgentId" TEXT NOT NULL,
    	"AgentName" TEXT NOT NULL,
    	"User" TEXT NOT NULL,
    	"Computer" TEXT NOT NULL,
    	"RemotePath" TEXT NOT NULL,
    	"LocalPath" TEXT NOT NULL,
    	"TotalSize" INTEGER,
    	"RecvSize" INTEGER,
    	"Date" BIGINT,
    	"State" INTEGER
    );`
	_, err = dbms.database.Exec(createTableQuery)

	createTableQuery = `CREATE TABLE IF NOT EXISTS "Screenshots" (
    	"ScreenId" TEXT NOT NULL UNIQUE, 
    	"User" TEXT NOT NULL,
    	"Computer" TEXT NOT NULL,
    	"LocalPath" TEXT NOT NULL,
    	"Note" TEXT,
    	"Date" BIGINT
    );`
	_, err = dbms.database.Exec(createTableQuery)

	createTableQuery = `CREATE TABLE IF NOT EXISTS "Agents" (
    	"Id" TEXT NOT NULL UNIQUE, 
    	"Crc" TEXT NOT NULL,
    	"Name" TEXT NOT NULL,
    	"SessionKey" BLOB NOT NULL,
    	"Listener" TEXT NOT NULL,
    	"Async" INTEGER,
    	"ExternalIP" TEXT,
    	"InternalIP" TEXT,
    	"GmtOffset" INTEGER,
    	"Sleep" INTEGER,
    	"Jitter" INTEGER,
    	"Pid" TEXT,
    	"Tid" TEXT,
    	"Arch" TEXT,
    	"Elevated" INTEGER,
    	"Process" TEXT,
    	"Os" INTEGER,
    	"OsDesc" TEXT,
    	"Domain" TEXT,
    	"Computer" TEXT,
    	"Username" TEXT,
    	"Impersonated" TEXT,
    	"OemCP" INTEGER,
    	"ACP" INTEGER,
    	"CreateTime" BIGINT,
    	"LastTick" INTEGER,
    	"WorkingTime" INTEGER,	
    	"KillDate" INTEGER,
    	"Tags" TEXT,
    	"Mark" TEXT,
    	"Color" TEXT,
    	"TargetId" TEXT,
    	"CustomData" BLOB
    );`
	_, err = dbms.database.Exec(createTableQuery)

	createTableQuery = `CREATE TABLE IF NOT EXISTS "Tasks" (
    	"TaskId" TEXT NOT NULL UNIQUE, 
    	"AgentId" TEXT NOT NULL,
    	"TaskType" INTEGER,
    	"Client" TEXT,
    	"User" TEXT,
    	"Computer" TEXT,
    	"StartDate" BIGINT,
    	"FinishDate" BIGINT,
    	"CommandLine" TEXT NOT NULL,
    	"MessageType" INTEGER,
    	"Message" TEXT,
    	"ClearText" TEXT,
    	"Completed" INTEGER
    );`
	_, err = dbms.database.Exec(createTableQuery)

	createTableQuery = `CREATE TABLE IF NOT EXISTS "Consoles" (
		"Id" INTEGER PRIMARY KEY AUTOINCREMENT,
    	"AgentId" TEXT NOT NULL,
    	"Packet" BLOB
    );`
	_, err = dbms.database.Exec(createTableQuery)

	createTableQuery = `CREATE TABLE IF NOT EXISTS "Pivots" (
		"Id" INTEGER PRIMARY KEY AUTOINCREMENT,
    	"PivotId" TEXT NOT NULL,
    	"PivotName" TEXT NOT NULL,
    	"ParentAgentId" TEXT NOT NULL,
    	"ChildAgentId" TEXT NOT NULL
    );`
	_, err = dbms.database.Exec(createTableQuery)

	createTableQuery = `CREATE TABLE IF NOT EXISTS "Credentials" (
		"Id" INTEGER PRIMARY KEY AUTOINCREMENT,
    	"CredId" TEXT NOT NULL,
    	"Username" TEXT,
    	"Password" TEXT,
    	"Realm" TEXT,
    	"Type" TEXT,
    	"Tag" TEXT,
    	"Date" BIGINT,
    	"Storage" TEXT,
		"AgentId" TEXT,
		"Host" TEXT
    );`
	_, err = dbms.database.Exec(createTableQuery)

	createTableQuery = `CREATE TABLE IF NOT EXISTS "Targets" (
		"Id" INTEGER PRIMARY KEY AUTOINCREMENT,
    	"TargetId" TEXT NOT NULL,
    	"Computer" TEXT,
    	"Domain" TEXT,
    	"Address" TEXT,
    	"Os" INTEGER,
    	"OsDesk" TEXT,
    	"Tag" TEXT,
    	"Info" TEXT,
    	"Date" BIGINT,
		"Alive" BOOLEAN,
		"Agents" TEXT
    );`
	_, err = dbms.database.Exec(createTableQuery)

	indexQueries := []string{
		`CREATE INDEX IF NOT EXISTS idx_tasks_agentid ON Tasks(AgentId);`,
		`CREATE INDEX IF NOT EXISTS idx_tasks_startdate ON Tasks(StartDate);`,
		`CREATE INDEX IF NOT EXISTS idx_tasks_completed ON Tasks(Completed);`,
		`CREATE INDEX IF NOT EXISTS idx_consoles_agentid ON Consoles(AgentId);`,
		`CREATE INDEX IF NOT EXISTS idx_downloads_agentid ON Downloads(AgentId);`,
		`CREATE INDEX IF NOT EXISTS idx_downloads_date ON Downloads(Date);`,
		`CREATE INDEX IF NOT EXISTS idx_screenshots_date ON Screenshots(Date);`,
		`CREATE INDEX IF NOT EXISTS idx_credentials_agentid ON Credentials(AgentId);`,
		`CREATE INDEX IF NOT EXISTS idx_pivots_parentagentid ON Pivots(ParentAgentId);`,
		`CREATE INDEX IF NOT EXISTS idx_pivots_childagentid ON Pivots(ChildAgentId);`,
	}
	for _, indexQuery := range indexQueries {
		_, _ = dbms.database.Exec(indexQuery)
	}

	return err
}

func (dbms *DBMS) DatabaseExists() bool {
	return dbms.exists
}