utils.fs

module Utils

open System
open System.IO
open System.Numerics
open System.Security.Cryptography

(* IO *)

let readLines (filePath: string) =
    seq { use sr = new StreamReader (filePath)
          while not sr.EndOfStream do
          yield sr.ReadLine () }

(* Helpers *)

let xor (b1: byte seq) (b2: byte seq) : byte seq =
    Seq.zip b1 b2
    |> Seq.map (fun (x, y) -> x ^^^ y)
    
let xorArr b1 b2 = xor b1 b2 |> Array.ofSeq

let xorArrs (lst: byte [] list) : byte [] =
    let rec f acc lst =
        match acc, lst with
            | [||], x::y::[] -> xorArr x y
            | [||], x::y::xs -> f (xorArr x y) xs
            | _,    x::[]    -> xorArr acc x
            | _,    x::xs    -> f (xorArr acc x) xs
    f [||] lst 

let rec transpose xss =
   match xss with
       | ([]::_) -> []
       | xss   -> List.map List.head xss :: transpose (List.map List.tail xss)

let repeat x = seq { while true do yield x }
let repeatArr x n = repeat x |> Seq.take n |> Seq.toArray
let repeatSeq xs = seq { while true do yield! xs }

(* Encodings *)

let hexToByte = function
    | '0' -> 0uy  | '1' -> 1uy
    | '2' -> 2uy  | '3' -> 3uy
    | '4' -> 4uy  | '5' -> 5uy
    | '6' -> 6uy  | '7' -> 7uy
    | '8' -> 8uy  | '9' -> 9uy
    | 'a' -> 10uy | 'b' -> 11uy
    | 'c' -> 12uy | 'd' -> 13uy
    | 'e' -> 14uy | 'f' -> 15uy
    | _ -> failwith "Invalid hex char"

let hexToBytes s =
    Seq.toList s
    |> List.chunkBySize 2 
    |> Seq.map (fun pair -> (Seq.head pair, Seq.last pair))
    |> Seq.map (fun (x, y) -> (hexToByte x <<< 4) ||| hexToByte y)
    |> List.ofSeq

let update i n b =
    ((b |> int |> BigInteger) <<< (i * 8)) ||| n

let bytesToBigInt bs =
    bs
    |> List.rev
    |> List.fold (fun (i, n) b -> (i + 1, update i n b)) (0, BigInteger 0)
    |> snd

let hexToBigInt s = hexToBytes s |> bytesToBigInt

let bytesToStr (b: byte seq) : string =
    Seq.toArray b
    |> Text.Encoding.ASCII.GetString
    
let bytesToHex (b: byte seq) : string =
    Seq.map (sprintf "%02x") b
    |> String.concat ""

let strToBytes (s: string) : byte [] =
    Text.Encoding.ASCII.GetBytes s
    
let histogram xs =
    Seq.groupBy id xs
    |> Map.ofSeq
    |> Map.map (fun k v -> Seq.length v)

(* PKCS7 Padding *)

let padPKCS7 (length: int) (bytes: byte seq) : byte [] =
    let targetLen = length - (Seq.length bytes % length)
    let padLen = if targetLen > 0 then targetLen else Seq.length bytes
    let pad = repeat (byte padLen)
    Seq.append bytes (Seq.take padLen pad) |> Seq.toArray

let validPKCS7 (code: byte []) : bool =
    let arr = Array.rev code
    let rec strip (arr: byte []) padding ctr : bool =
        let c = int arr.[0]
        if c <> padding then false
        elif ctr < padding then strip arr.[1..] padding (ctr + 1)
        else true
    match arr.[0] with
        | 0uy -> false
        | _   -> strip arr (int arr.[0]) 1

let stripPKCS7 (code: byte []) : byte [] =
    let arr = Array.rev code
    let rec strip (arr: byte []) padding ctr : byte [] =
        let c = int arr.[0]
        if c <> padding then failwith "bad padding"
        elif ctr < padding then strip arr.[1..] padding (ctr + 1)
        else arr.[1..] |> Array.rev
    strip arr (int arr.[0]) 1
        
(* AES *)

let randKey (size: int) : byte [] =
    let rnd = Random()
    [|for _ in 1..size do yield rnd.Next 256 |> byte|]

let genAES (key: string) =
    let aes = new AesManaged()
    aes.Mode <- CipherMode.ECB
    aes.Key <- strToBytes key
    aes.Padding <- PaddingMode.None
    aes

let prepareInputECB (input: byte []) : byte [] [] =
    input
    |> padPKCS7 16
    |> Array.chunkBySize 16

let applyAESEncryptECB (key: string) (input: byte [] []) : byte [] =
    use aes = genAES key
    let encryptor = aes.CreateEncryptor(aes.Key, aes.IV)
    [| for block in input do
           let encrypted = Array.create 16 0uy
           encryptor.TransformBlock(block, 0, 16, encrypted, 0) |> ignore
           yield! encrypted |]

let applyAESDecryptECB (key: string) (code: byte []) : byte [] =
    let aes = genAES key
    let decryptor = aes.CreateDecryptor(aes.Key, aes.IV)
    let codeLen = Array.length code
    let decrypted = Array.create codeLen 0uy
    decryptor.TransformBlock(code, 0, codeLen, decrypted, 0) |> ignore
    decrypted

let AESEncryptECB (key: string) (iv: byte []) (input: byte []) : byte [] =
    applyAESEncryptECB key (prepareInputECB input)

let AESDecryptECB (key: string) (code: byte []) : byte [] =
    applyAESDecryptECB key code
    
(* CBC *)

let IV = Seq.take 16 (repeat (byte 0)) |> Seq.toArray

let prepareInputCBC (input: byte []) : byte [] list  =
    input
    |> padPKCS7 16
    |> Array.chunkBySize 16
    |> Array.toList

let prepareCodeCBC(code: byte []) : byte [] list =
    code
    |> Array.chunkBySize 16
    |> Array.toList

let rec applyCBCEncrypt blocks key acc : byte [] list =
    let genArray = Seq.toArray >> Array.create 1
    let encrypt x y = xor x y |> genArray |> applyAESEncryptECB key
    match blocks with
        | x::y::[] -> let encrypted = encrypt x y
                      List.rev (encrypted::acc)
        | x::y::xs -> let encrypted = encrypt x y
                      applyCBCEncrypt (encrypted::xs) key (encrypted::acc)

let rec applyCBCDecrypt blocks key acc : byte [] list =
    let decrypt x y = applyAESDecryptECB key y |> Array.ofSeq |> xor x |> Seq.toArray
    match blocks with
        | x::y::[] -> let decrypted = decrypt x y
                      List.rev (decrypted::acc)
        | x::y::xs -> let decrypted = decrypt x y
                      applyCBCDecrypt (y::xs) key (decrypted::acc)

let CBCEncrypt (key: string) (iv: byte []) (input: byte []) : byte [] = 
    let blocks = iv :: (prepareInputCBC input)
    applyCBCEncrypt blocks key [] |> Array.concat
   
let CBCDecrypt (key: string) (iv: byte []) (code: byte []) : byte [] =
    let blocks = iv :: (prepareCodeCBC code)
    applyCBCDecrypt blocks key [] |> Array.concat |> stripPKCS7

let CBCDecryptKeepPad (key: string) (iv: byte []) (code: byte []) : byte [] =
    let blocks = iv :: (prepareCodeCBC code)
    applyCBCDecrypt blocks key [] |> Array.concat

(* RSA *)

let rec egcd (a: int) (b: int) : (int * int * int) =
    match a, b with
    | 0, b -> (b, 0, 1)
    | a, b -> let (g, s, t) = egcd (b % a) a in
              (g, (t - (b / a) * s), s)

let modInv a m : int option =
    let g, s, _ = egcd a m
    let mkPos n = if n < 0 then n + m else n
    if g = 1 then Some (mkPos s) else None

let rec egcdBig (a: BigInteger) (b: BigInteger) : (BigInteger * BigInteger * BigInteger) =
    if a = (BigInteger 0) then (b, (BigInteger 0), (BigInteger 1))
    else let (g, s, t) = egcdBig (b % a) a in (g, (t - (b / a) * s), s)

let modInvBig a m : BigInteger option =
    let g, s, _ = egcdBig a m
    let mkPos n = if n < (BigInteger 0) then n + m else n
    if g = (BigInteger 1) then Some (mkPos s) else None

let primes = readLines "large_primes.csv" |> Seq.map BigInteger.Parse

let genRSAKeys (r: Random) : ((BigInteger * BigInteger) * (BigInteger * BigInteger)) =
    let rec pick e =
        let p = Seq.take (1 + r.Next(9999)) primes |> Seq.last
        if p % e = (BigInteger 0) then pick e else p
    let e = BigInteger 3
    let p = pick e
    let q = pick e
    let n = p * q
    let et = (p - (BigInteger 1)) * (q - (BigInteger 1))
    let d = (modInvBig e et).Value
    ((e, n), (d, n))