package main

import (
    "os"
    "fmt"
    "bufio"
    "config"
    "regexp"
    "strings"
    "tools/rsa"
    "tools/rsa/data"
    "tools/common"
)

type optval int
const (
    STD optval = 1 << iota
    GEN optval = 1 << iota
    RAGNAR optval = 1 << iota
    BIRGER optval = 1 << iota
    MATHEM optval = 1 << iota
    CUSTOM optval = 1 << iota
    SIGN optval = 1 << iota
    CRYPTSIGN optval = 1 << iota
    EVERY optval = CUSTOM | SIGN | CRYPTSIGN
)


func opt() optval {

    fmt.Printf("Select option\n")
    fmt.Printf("1: Set standard keys\n")
    fmt.Printf("2: Generate new keys\n")
    s := userInput("Option", false)

    switch s[0] {
        case '1':
            return STD
        case '2':
            return GEN
        default:
            fmt.Printf("Unknown option, set standard.\n")
            return STD
    }
}

func cryptOpt() optval {

    fmt.Printf("Select option\n")
    fmt.Printf("1: Read ragnars message\n")
    fmt.Printf("2: Send Birger a message\n")
    fmt.Printf("3: Encrypt/decrypt mathematica message\n")
    fmt.Printf("4: Encrypt/decrypt custom message\n")
    fmt.Printf("5: Sign/de-sign custom message\n")
    fmt.Printf("6: Encrypt+sign/decrypte+unsign custom message\n")
    fmt.Printf("7: Encrypt, sign and encrypt+sign\n")
    s := userInput("Option", false)

    switch s[0] {
        case '1': return RAGNAR
        case '2': return BIRGER
        case '3': return MATHEM
        case '4': return CUSTOM
        case '5': return SIGN
        case '6': return CRYPTSIGN
        case '7': return EVERY
        default:
            fmt.Printf("Unknown option, set standard.\n")
            return MATHEM
    }
}

func std1(r *rsa.RSA) bool {
    fmt.Printf("* Setting standard-keys for r1\n")
    return r.SetKeys(config.RSA_STD_1_D, config.RSA_STD_1_E, config.RSA_STD_1_N)
}

func std2(r *rsa.RSA) bool {
    fmt.Printf("* Setting standard-keys for r2\n")
    return r.SetKeys(config.RSA_STD_2_D, config.RSA_STD_2_E, config.RSA_STD_2_N)
}

func stdRagnar(r *rsa.RSA) bool {
    fmt.Printf("* Setting ragnar-keys for r2\n")
    return r.SetKeys(config.RSA_RAGNAR_D, config.RSA_RAGNAR_E, config.RSA_RAGNAR_N)
}

func stdBirger(r *rsa.RSA) bool {
    fmt.Printf("* Setting ragnar-keys for r2\n")
    return r.SetKeys(config.RSA_BIRGER_D, config.RSA_BIRGER_E, config.RSA_BIRGER_N)
}

func std(o optval, r1, r2 *rsa.RSA) bool {
    switch o {
        case RAGNAR:
            return std1(r1) && stdRagnar(r2)
        case BIRGER:
            return std1(r1) && stdBirger(r2)
        default:
            return std1(r1) && std2(r2)
    }
}

func gen(o optval, r1, r2 *rsa.RSA) bool {
    switch o {
        case RAGNAR, EVERY:
            r1.GenPrimes()
            stdRagnar(r2)
        case BIRGER:
            r1.GenPrimes()
            stdBirger(r2)
        default:
            r1.GenPrimes()
            r2.GenPrimes()
    }
    return true
}

func userInput(promt string, re bool) string {
    reader := bufio.NewReader(os.Stdin)
    fmt.Printf("%s [enter]: ", promt)
    msg, err := reader.ReadString('\n')
    if !common.IsError(err, "Could not read from input\n") {
        if re {
            s := regexp.MustCompile("[a-zæøåA-ZÆØÅ., ]+").FindAllString(msg, -1) // Allowin a-z, æ, ø, å, comma, dot, space
            msg = ""
            for i := 0; i < len(s); i++ {
                msg += strings.ToLower(strings.Trim(s[i], "\n"))
            }
        }
        return msg
    }
    return ""
}

func main() {

    r1 := rsa.New()
    r2 := rsa.New()

    if r1 == nil || r2 == nil {
        fmt.Printf("Could not create RSA")
        return
    }

    fmt.Printf("R1: %p\tR2: %p\n", r1, r2)

    var keys bool
    m := cryptOpt()
    o := opt()
    switch o {
        case STD:
            keys = std(m, r1, r2)
        case GEN:
            keys = gen(m, r1, r2)
    }

    if !keys {
        fmt.Printf("Oh no!")
        return
    }

    switch m {
        case RAGNAR:
            r2.Crack()
            msg,_ := r2.ReadRecvFile()
            d := r2.Set(data.TXT, data.Data(""))
            d.Data().SetRaw(msg)
            enc := d.AsNumbers().ToList()
            fmt.Printf("%s\n", d.Repr())
            dec := r2.Decrypt(enc)
            fmt.Printf("%s\n", dec)
        case BIRGER:
            msg := userInput("Message to Birger", true)
            enc := r1.Encrypt(msg, r2.Public().N, r2.Public().E)
            fmt.Printf("Message to send to Birger:\n%v\n\n", enc)
        case MATHEM:
            msg := "dette er en liten test på hva som går an å gjøre med mathematica, mye annet og naturligvis."
            fmt.Printf("Message to encrypt:\n%s\n\n", msg)
            enc := r1.Encrypt(msg, r2.Public().N, r2.Public().E)
            fmt.Printf("Encrypted message:\n%s\n\n", enc)
            dec := r2.Decrypt(enc)
            fmt.Printf("Decryptet message:\n%s\n\n", dec)
        case CUSTOM:
            msg := userInput("Your message", true)
            fmt.Printf("Message to encrypt:\n%s\n\n", msg)
            enc := r1.Encrypt(msg, r2.Public().N, r2.Public().E)
            fmt.Printf("Encrypted:\n%v\n\n", enc)
            dec := r2.Decrypt(enc)
            fmt.Printf("DECRYPTED MESSAGE:\n%s\n", dec)
        case SIGN:
            msg := userInput("Your message", true)
            fmt.Printf("Message to sign:\n%s\n\n", msg)
            signed := r1.Sign(msg)
            fmt.Printf("Signed:\n%v\n\n", signed)
            unsign := r1.Unsign(signed, r1.Public().N, r1.Public().E)
            fmt.Printf("Signed message:\n%s\n\n", unsign)
        case CRYPTSIGN:
            msg := userInput("Your message", true)
            fmt.Printf("Message to crypt and sign:\n%s\n\n", msg)
            encsign := r1.EncryptNSign(msg, r2.Public().N, r2.Public().E)
            fmt.Printf("Encrypted and signed:\n%v\n\n", encsign)
            decunsign := r2.DecryptNUnsign(encsign, r1.Public().N, r1.Public().E)
            fmt.Printf("Uncrypted and unsigned msg:\n%s\n\n", decunsign)
        case EVERY:
            msg := userInput("Your message", true)
            fmt.Printf("Message:\n%s\n\n", msg)
            enc := r1.Encrypt(msg, r2.Public().N, r2.Public().E)
            signed := r1.Sign(msg)
            encsign := r1.EncryptNSign(msg, r2.Public().N, r2.Public().E)
            fmt.Printf("Encrypted:\n%v\n\nSigned:\n%v\n\nEncrypted and signed:\n%s\n\n", enc, signed, encsign)
    }
}
