from assemblyline import odm
from assemblyline.common.attack_map import attack_map
from assemblyline.odm.models.ontology.results.network import REQUEST_METHODS, LOOKUP_TYPES

# Based on model in MaCo framework
CATEGORIES = ["adware", "apt", "backdoor", "banker", "bootkit", "bot", "browser_hijacker", "bruteforcer",
              "clickfraud", "cryptominer", "ddos", "downloader", "dropper", "exploitkit", "fakeav", "hacktool",
              "infostealer", "keylogger", "loader", "obfuscator", "pos", "proxy", "rat", "ransomware",
              "reverse_proxy", "rootkit", "scanner", "scareware", "spammer", "trojan", "virus", "wiper",
              "webshell", "worm"]
CONNECTION_USAGE = ["c2", "upload", "download", "propagate", "tunnel", "other", "ransom", "decoy"]


@odm.model(description="Encryption details")
class Encryption(odm.Model):
    algorithm = odm.Optional(odm.Text(), description="Algorithm")
    public_key = odm.Optional(odm.Text(), description="Public Key")
    key = odm.Optional(odm.Text(), description="Key")
    provider = odm.Optional(odm.Text(), description="Provider")
    mode = odm.Optional(odm.Text(), description="Mode")
    iv = odm.Optional(odm.Text(), description="Initialization Vector")
    seed = odm.Optional(odm.Text(), description="Seed")
    nonce = odm.Optional(odm.Text(), description="Nonce value")
    constants = odm.Optional(odm.List(odm.Text()), description="Constants")
    usage = odm.Optional(odm.Enum(values=["config", "communication", "binary", "ransom", "other"]),
                         description="Purpose of encryptions")


@odm.model(description="Binary data extracted by decoder")
class Binary(odm.Model):
    datatype = odm.Optional(odm.Enum(values=["payload", "config", "other"]))
    data = odm.Text()
    other = odm.Optional(odm.Mapping(odm.Any()), description="Other information")
    encryption = odm.Optional(odm.List(odm.Compound(Encryption)))


@odm.model(description="Usage of FTP connection")
class FTP(odm.Model):
    username = odm.Optional(odm.Text(), description="Username")
    password = odm.Optional(odm.Text(), description="Password")
    hostname = odm.Optional(odm.Text(), description="FTP Host")
    port = odm.Optional(odm.Integer(), description="FTP Port")
    path = odm.Optional(odm.Text(), description="FTP Path")
    usage = odm.Optional(odm.Enum(values=CONNECTION_USAGE), description="Purpose of FTP connection")


@odm.model(description="Usage of SMTP connection")
class SMTP(odm.Model):
    username = odm.Optional(odm.Text(), description="Username")
    password = odm.Optional(odm.Text(), description="Password")
    hostname = odm.Optional(odm.Text(), description="SMTP Host")
    port = odm.Optional(odm.Integer(), description="SMTP Port")
    mail_to = odm.Optional(odm.List(odm.Text()), description="Sent to")
    mail_from = odm.Optional(odm.Text(), description="Sent from")
    subject = odm.Optional(odm.Text(), description="Subject")
    usage = odm.Optional(odm.Enum(values=CONNECTION_USAGE), description="Purpose of SMTP connection")


@odm.model(description="Usage of HTTP connection")
class HTTP(odm.Model):
    uri = odm.Optional(odm.URI(), description="URI")
    protocol = odm.Optional(odm.Enum(values=["http", "https"]), description="Protocol")
    username = odm.Optional(odm.Text(), description="Username")
    password = odm.Optional(odm.Text(), description="Password")
    hostname = odm.Optional(odm.Text(), description="HTTP server")
    port = odm.Optional(odm.Integer(), description="HTTP Port")
    path = odm.Optional(odm.URIPath(), description="URI Path")
    query = odm.Optional(odm.Text(), description="Query parameters")
    fragment = odm.Optional(odm.Text(), description="Fragment")
    user_agent = odm.Optional(odm.Text(), description="User Agent")
    method = odm.Optional(odm.Enum(values=REQUEST_METHODS), description="Method")
    headers = odm.Optional(odm.Mapping(odm.Text()), description="HTTP Headers")
    max_size = odm.Optional(odm.Integer(), description="Maximum size")
    usage = odm.Optional(odm.Enum(values=CONNECTION_USAGE), description="Purpose of HTTP connection")


@odm.model(description="Usage of SSH connection")
class SSH(odm.Model):
    username = odm.Optional(odm.Text(), description="Username")
    password = odm.Optional(odm.Text(), description="Password")
    public_key = odm.Optional(odm.Text(), description="SSH Public Key")
    hostname = odm.Optional(odm.Text(), description="SSH Host")
    port = odm.Optional(odm.Integer(), description="SSH Port")
    usage = odm.Optional(odm.Enum(values=CONNECTION_USAGE), description="Purpose of SSH connection")


@odm.model(description="Usage of Proxy connection")
class Proxy(odm.Model):
    username = odm.Optional(odm.Text(), description="Username")
    password = odm.Optional(odm.Text(), description="Password")
    hostname = odm.Optional(odm.Text(), description="Proxy Host")
    port = odm.Optional(odm.Integer(), description="Proxy Port")
    usage = odm.Optional(odm.Enum(values=CONNECTION_USAGE), description="Purpose of proxy connection")
    protocol = odm.Optional(odm.Text(), description="Protocol used")


@odm.model(description="Usage of DNS connection")
class DNS(odm.Model):
    ip = odm.Optional(odm.IP(), description="IP of DNS server")
    port = odm.Optional(odm.Integer(), description="Port of DNS server")
    hostname = odm.Optional(odm.Text(), description="Hostname used in query")
    record_type = odm.Optional(odm.Enum(values=LOOKUP_TYPES), description="Type of DNS record")
    usage = odm.Optional(odm.Enum(values=CONNECTION_USAGE), description="Purpose of DNS connection")


@odm.model(description="Usage of General TCP/UDP connection")
class GeneralConnection(odm.Model):
    client_ip = odm.Optional(odm.IP(), description="Client IP")
    client_port = odm.Optional(odm.Integer(), description="Client Port")
    server_ip = odm.Optional(odm.IP(), description="Server IP")
    server_domain = odm.Optional(odm.Domain(), description="Server Domain")
    server_port = odm.Optional(odm.Integer(), description="Server Port")
    usage = odm.Optional(odm.Enum(values=CONNECTION_USAGE), description="Purpose of TCP/UDP connection")


@odm.model(description="Operating System services affected")
class Service(odm.Model):
    dll = odm.Optional(odm.Text(), description="DLL associated to service")
    name = odm.Optional(odm.Text(), description="Name of service")
    display_name = odm.Optional(odm.Text(), description="Display Name of service")
    description = odm.Optional(odm.Text(), description="Service Description")


@odm.model(description="Cryptocoin usage (ransomware/miner)")
class Cryptocurrency(odm.Model):
    coin = odm.Optional(odm.Text(), description="Name of coin used")
    address = odm.Optional(odm.Text(), description="Wallet address")
    random_amount = odm.Optional(odm.Integer(), description="Ransom amount")
    usage = odm.Optional(odm.Enum(values=['ransomware', 'miner', 'other']), description="Use of cryptocurrency")


@odm.model(description="File Paths")
class Path(odm.Model):
    path = odm.Optional(odm.Text(), description="Path")
    usage = odm.Optional(
        odm.Enum(values=['c2', 'config', 'install', 'plugins', 'logs', 'storage', 'other']),
        description="Use of path")


@odm.model(description="Inter-Process Communications")
class IPC(odm.Model):
    file = odm.Optional(odm.List(odm.Text()),
                        description="A record stored on disk, or a record synthesized on demand by a file server, "
                        "which can be accessed by multiple processes.")
    socket = odm.Optional(odm.List(odm.Text()),
                          description="Data sent over a network interface, either to a different process on the same "
                          "computer or to another computer on the network. Stream oriented (TCP; data written through a"
                          " socket requires formatting to preserve message boundaries) or more rarely message-oriented "
                          "(UDP, SCTP).")
    unix_domain_socket = odm.Optional(odm.List(odm.Text()),
                                      description="Similar to an internet socket, but all communication occurs within "
                                      "the kernel. Domain sockets use the file system as their address space. "
                                      "Processes reference a domain socket as an inode, and multiple processes can "
                                      "communicate with one socket.")
    memory_mapped_file = odm.Optional(odm.List(odm.Text()),
                                      description="A file mapped to RAM and can be modified by changing memory"
                                      "addresses directly instead of outputting to a stream. This shares the same "
                                      "benefits as a standard file.")
    message_queue = odm.Optional(odm.List(odm.Text()),
                                 description="A data stream similar to a socket, but which usually preserves message "
                                 "boundaries. Typically implemented by the operating system, they allow multiple "
                                 "processes to read and write to the message queue without being directly connected to "
                                 "each other.")
    anonymous_pipe = odm.Optional(odm.List(odm.Text()),
                                  description="A unidirectional data channel using standard input and output. Data "
                                  "written to the write-end of the pipe is buffered by the operating system until it "
                                  "is read from the read-end of the pipe. Two-way communication between processes can "
                                  'be achieved by using two pipes in opposite "directions".')
    named_pipe = odm.Optional(odm.List(odm.Text()),
                              description="A pipe that is treated like a file. Instead of using standard input and "
                              "output as with an anonymous pipe, processes write to and read from a named pipe, as if "
                              "it were a regular file.")
    process_names = odm.Optional(odm.List(odm.Text()),
                                 description="The process names involved in the IPC communication")
    shared_memory = odm.Optional(odm.Text(),
                                 description="Multiple processes are given access to the same block of memory, which "
                                 "creates a shared buffer for the processes to communicate with each other.")
    usage = odm.Optional(odm.Enum(values=CONNECTION_USAGE), description="Purpose of connection")


@odm.model(description="Registry")
class Registry(odm.Model):
    key = odm.Text(description="Registry key")
    value = odm.Optional(odm.Text(), description="Registry")
    usage = odm.Optional(
        odm.Enum(values=['persistence', 'store_data', 'store_payload', 'read', 'other']),
        description="Use of registry key")


@odm.model(description="Extracted Malware Configuration")
class MalwareConfig(odm.Model):
    config_extractor = odm.Keyword(description="Name of extractor")

    family = odm.List(odm.Text(), description="What family is this associated to?")
    version = odm.Optional(odm.Text(), description="Version of the malware")
    category = odm.Optional(odm.List(odm.Enum(values=CATEGORIES)), description="Category of malware")

    attack = odm.Optional(odm.List(odm.Enum(values=list(attack_map.keys()))), description="ATT&CK ID associated")
    capability_enabled = odm.Optional(odm.List(odm.Text()), description="Enabled Capabilities")
    capability_disabled = odm.Optional(odm.List(odm.Text()), description="Disabled Capabilities")
    campaign_id = odm.Optional(odm.List(odm.Text()), description="Campaign ID")
    identifier = odm.Optional(odm.List(odm.Text()), description="Identifier")
    decoded_strings = odm.Optional(odm.List(odm.Text()), description="Decoded Strings")
    password = odm.Optional(odm.List(odm.Text()), description="Passwords")
    mutex = odm.Optional(odm.List(odm.Text()), description="Mutex")
    pipe = odm.Optional(odm.List(odm.Text()), description="Pipe")
    ipc = odm.Optional(odm.List(odm.Compound(IPC)), description="IPC (similar to 'pipe' field but more detailed)")
    sleep_delay = odm.Optional(odm.Integer(), description="Sleep Delay")
    sleep_delay_jitter = odm.Optional(odm.Integer(), description="Sleep Delay Jitter")
    inject_exe = odm.Optional(odm.List(odm.Text()), description="Injected EXE")

    binaries = odm.Optional(odm.List(odm.Compound(Binary)), description="Binaries")
    ftp = odm.Optional(odm.List(odm.Compound(FTP)), description="FTPs")
    smtp = odm.Optional(odm.List(odm.Compound(SMTP)), description="SMTPs")
    http = odm.Optional(odm.List(odm.Compound(HTTP)), description="HTTPs")
    ssh = odm.Optional(odm.List(odm.Compound(SSH)), description="SSHs")
    proxy = odm.Optional(odm.List(odm.Compound(Proxy)), description="Proxies")
    dns = odm.Optional(odm.List(odm.Compound(DNS)), description="DNS")
    tcp = odm.Optional(odm.List(odm.Compound(GeneralConnection)), description="TCPs")
    udp = odm.Optional(odm.List(odm.Compound(GeneralConnection)), description="UDPs")
    encryption = odm.Optional(odm.List(odm.Compound(Encryption)), description="Encryptions")
    service = odm.Optional(odm.List(odm.Compound(Service)), description="Services")
    cryptocurrency = odm.Optional(odm.List(odm.Compound(Cryptocurrency)), description="Cryptocurrencies")
    paths = odm.Optional(odm.List(odm.Compound(Path)), description="Paths")
    registry = odm.Optional(odm.List(odm.Compound(Registry)), description="Registry")

    other = odm.Optional(odm.Mapping(odm.Any()), description="Other information")
