<commit_before>import pandas as pd
from datetime import datetime

from poshc2.Colours import Colours
from poshc2.server.database.Model import C2, Implant, NewTask, HostedFile
from poshc2.server.database.DBType import DBType
from poshc2.server.Config import Database, DatabaseType, PoshProjectDirectory


if DatabaseType == DBType.Postgres:
    import poshc2.server.database.DBPostgres as DBImplementation
else:
    import poshc2.server.database.DBSQLite as DBImplementation


conn = None


def get_conn():
    global conn
    return conn


def set_conn(new_conn):
    global conn
    conn = new_conn


def database_connect():
    if conn is None:
        set_conn(DBImplementation.connect())


def initializedb():
    database_connect()
    DBImplementation.initialise(create_database)


def db_exists():
    if get_conn() is None:
        database_connect()
    return DBImplementation.db_exists(conn)


def create_database(create_urls, create_implants, create_autoruns, create_tasks, create_newtasks,
                    create_creds, create_opsec_entry, create_c2server, create_c2_messages, create_hosted_files,
                    create_power_status):
    try:
        c = get_conn().cursor()
    except Exception as e:
        print("[-] Error occurred using %s" % Database)
        print("[-] Exception: %s" % e)

    if conn is not None:
        try:
            c.execute(create_urls)
            c.execute(create_implants)
            c.execute(create_autoruns)
            c.execute(create_tasks)
            c.execute(create_newtasks)
            c.execute(create_creds)
            c.execute(create_opsec_entry)
            c.execute(create_c2server)
            c.execute(create_c2_messages)
            c.execute(create_hosted_files)
            c.execute(create_power_status)
            conn.commit()
        except Exception as e:
            print("Error creating database: " + str(e))


def convert_query(query, postgres_suffix=None):
    if DatabaseType == DBType.Postgres:
        query = query.replace("?", "%s")
        if postgres_suffix:
            query += postgres_suffix
    return query


def get_last_insert_row_id(cursor):
    if DatabaseType == DBType.Postgres:
        return cursor.fetchone()[0]
    return cursor.lastrowid


def setupserver(PayloadCommsHost, EncKey, DomainFrontHeader, DefaultSleep, KillDate, GET_404_Response, PoshProjectDirectory, QuickCommand, DownloadURI, ProxyURL, ProxyUser, ProxyPass, URLS, SocksURLS, Insecure, UserAgent, Referrer, Pushover_APIToken, Pushover_APIUser, Slack_UserID, Slack_Channel, Slack_BotToken, EnableNotifications):
    c = conn.cursor()
    command = convert_query("INSERT INTO C2Server (PayloadCommsHost,EncKey,DomainFrontHeader,DefaultSleep,KillDate,GET_404_Response,PoshProjectDirectory,QuickCommand,DownloadURI,ProxyURL,ProxyUser,ProxyPass,URLS,SocksURLS,Insecure,UserAgent,Referrer,Pushover_APIToken,Pushover_APIUser,Slack_UserID,Slack_Channel,Slack_BotToken,EnableNotifications) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)")
    c.execute(command, (PayloadCommsHost, EncKey, DomainFrontHeader, DefaultSleep, KillDate, GET_404_Response, PoshProjectDirectory, QuickCommand, DownloadURI, ProxyURL, ProxyUser, ProxyPass, URLS, SocksURLS, Insecure, UserAgent, Referrer, Pushover_APIToken, Pushover_APIUser, Slack_UserID, Slack_Channel, Slack_BotToken, EnableNotifications))
    conn.commit()


def get_c2server_all():
    c = get_conn().cursor()
    c.execute("SELECT * FROM C2Server")
    result = c.fetchone()
    return C2(result[1], result[2], result[3], result[4], result[5], result[6], result[7], result[8], result[9],
              result[10], result[11], result[12], result[13], result[14], result[15], result[16], result[17],
              result[18], result[19], result[20])


def get_implants_all():
    c = get_conn().cursor()
    c.execute("SELECT * FROM Implants")
    results = c.fetchall()
    implants = []
    for result in results:
        implants.append(Implant(result[0], result[1], result[2], result[3], result[4], result[5], result[6], result[7], result[8],
                                result[9], result[10], result[11], result[12], result[13], result[14], result[15], result[16]))
    return implants

<commit_msg>select * needed number of columns returned updating<commit_after>import pandas as pd
from datetime import datetime

from poshc2.Colours import Colours
from poshc2.server.database.Model import C2, Implant, NewTask, HostedFile
from poshc2.server.database.DBType import DBType
from poshc2.server.Config import Database, DatabaseType, PoshProjectDirectory


if DatabaseType == DBType.Postgres:
    import poshc2.server.database.DBPostgres as DBImplementation
else:
    import poshc2.server.database.DBSQLite as DBImplementation


conn = None


def get_conn():
    global conn
    return conn


def set_conn(new_conn):
    global conn
    conn = new_conn


def database_connect():
    if conn is None:
        set_conn(DBImplementation.connect())


def initializedb():
    database_connect()
    DBImplementation.initialise(create_database)


def db_exists():
    if get_conn() is None:
        database_connect()
    return DBImplementation.db_exists(conn)


def create_database(create_urls, create_implants, create_autoruns, create_tasks, create_newtasks,
                    create_creds, create_opsec_entry, create_c2server, create_c2_messages, create_hosted_files,
                    create_power_status):
    try:
        c = get_conn().cursor()
    except Exception as e:
        print("[-] Error occurred using %s" % Database)
        print("[-] Exception: %s" % e)

    if conn is not None:
        try:
            c.execute(create_urls)
            c.execute(create_implants)
            c.execute(create_autoruns)
            c.execute(create_tasks)
            c.execute(create_newtasks)
            c.execute(create_creds)
            c.execute(create_opsec_entry)
            c.execute(create_c2server)
            c.execute(create_c2_messages)
            c.execute(create_hosted_files)
            c.execute(create_power_status)
            conn.commit()
        except Exception as e:
            print("Error creating database: " + str(e))


def convert_query(query, postgres_suffix=None):
    if DatabaseType == DBType.Postgres:
        query = query.replace("?", "%s")
        if postgres_suffix:
            query += postgres_suffix
    return query


def get_last_insert_row_id(cursor):
    if DatabaseType == DBType.Postgres:
        return cursor.fetchone()[0]
    return cursor.lastrowid


def setupserver(PayloadCommsHost, EncKey, DomainFrontHeader, DefaultSleep, KillDate, GET_404_Response, PoshProjectDirectory, QuickCommand, DownloadURI, ProxyURL, ProxyUser, ProxyPass, URLS, SocksURLS, Insecure, UserAgent, Referrer, Pushover_APIToken, Pushover_APIUser, Slack_UserID, Slack_Channel, Slack_BotToken, EnableNotifications):
    c = conn.cursor()
    command = convert_query("INSERT INTO C2Server (PayloadCommsHost,EncKey,DomainFrontHeader,DefaultSleep,KillDate,GET_404_Response,PoshProjectDirectory,QuickCommand,DownloadURI,ProxyURL,ProxyUser,ProxyPass,URLS,SocksURLS,Insecure,UserAgent,Referrer,Pushover_APIToken,Pushover_APIUser,Slack_UserID,Slack_Channel,Slack_BotToken,EnableNotifications) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)")
    c.execute(command, (PayloadCommsHost, EncKey, DomainFrontHeader, DefaultSleep, KillDate, GET_404_Response, PoshProjectDirectory, QuickCommand, DownloadURI, ProxyURL, ProxyUser, ProxyPass, URLS, SocksURLS, Insecure, UserAgent, Referrer, Pushover_APIToken, Pushover_APIUser, Slack_UserID, Slack_Channel, Slack_BotToken, EnableNotifications))
    conn.commit()


def get_c2server_all():
    c = get_conn().cursor()
    c.execute("SELECT * FROM C2Server")
    result = c.fetchone()
    return C2(result[1], result[2], result[3], result[4], result[5], result[6], result[7], result[8], result[9],
              result[10], result[11], result[12], result[13], result[14], result[15], result[16], result[17],
              result[18], result[19], result[20], result[21], result[22], result[23])


def get_implants_all():
    c = get_conn().cursor()
    c.execute("SELECT * FROM Implants")
    results = c.fetchall()
    implants = []
    for result in results:
        implants.append(Implant(result[0], result[1], result[2], result[3], result[4], result[5], result[6], result[7], result[8],
                                result[9], result[10], result[11], result[12], result[13], result[14], result[15], result[16]))
    return implants


{'fraction_non_alphanumeric': 0.11310216430067489, 'fraction_numerical': 0.01931580172213172, 'mean_word_length': 3.8477157360406093, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '788994', 'n_tokens_mistral': 2777, 'n_tokens_neox': 2668, 'n_words': 557}