golang使用SQLite数据库对指定扩展名文件进行批量加密（七、完整代码）

package main

import (
	"bufio"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"strings"
	"sync"
	"time"

	_ "github.com/mattn/go-sqlite3"
)

// 加密字符串
func GetNegation(strb []byte) []byte {
	var strsN []byte
	for i, b := range strb {
		ii := i ^ (i % 3)
		strsN = append(strsN, ^b-uint8(ii))
	}
	return strsN
}

//md5值,参数为字符串，输出md5字符串
func GetSha256(str string) string {
	strM := sha256.Sum256([]byte(str))
	return hex.EncodeToString(strM[:])
}

// 保存修改后的文件，文件名为md5值，内容为1024字节后内容
func SaveFile(fileMd5 string, content []byte) bool {
	contentNeg := GetNegation(content)
	err := ioutil.WriteFile(fileMd5, contentNeg, 0777) //将源文件，写入目标文件
	if err != nil {
		return false
	}
	return true
}

//返回文件名列表
func ReadFolder(fileDir string, extlist []string) []string {
	var fileNameList []string
	files, _ := ioutil.ReadDir(fileDir) //读取目录
	for _, onefile := range files {     //遍历目录下文件
		if !onefile.IsDir() { //是文件
			fileName := onefile.Name()
			fileNames := strings.Split(fileName, ".")
			ext := fileNames[len(fileNames)-1]
			for _, e := range extlist {
				if strings.EqualFold(ext, e) {
					fileNameList = append(fileNameList, fileName)
					break
				}
			}
		}
	}
	return fileNameList
}

//返回加密后的文件列表
func ReadFolderNeg(fileDir string) []string {
	var fileNameList []string
	files, _ := ioutil.ReadDir(fileDir) //读取目录
	for _, onefile := range files {     //遍历目录下文件
		if !onefile.IsDir() { //是文件
			fileName := onefile.Name()
			if len(fileName) == 64 {
				fileNameList = append(fileNameList, fileName)
			}
		}
	}
	return fileNameList
}

//删除文件
func DelFile(file *os.File, path string, jiajie string) {
	if file != nil {
		file.Close()
	}
	if IsFileExist(path) {
		err := os.Remove(path)
		if err == nil {
			// fmt.Printf("%s -> %s删除成功！\n", path, jiajie)
			return
		}
	}
	// fmt.Printf("%s -> %s删除失败！\n", path, jiajie)
}

//判断文件是否存在
func IsFileExist(path string) bool {
	fileInfo, err := os.Stat(path)
	if err != nil {
		return false
	}
	if fileInfo.IsDir() { //是目录
		return false
	}
	return true
}

//错误处理
func CheckErr(err error) {
	if err != nil {
		panic(err)
	}
}

//保存到数据库
func SaveFileToDb(db *sql.DB, fileName256 string, fileName string, content string, password string) bool {
	stmt, err := db.Prepare("insert into secret_key(fileName256,fileName,content,password) values(?,?,?,?)")
	defer func() {
		if stmt != nil {
			stmt.Close()
		}
	}()
	CheckErr(err)
	rwm.Lock()
	res, err := stmt.Exec(fileName256, fileName, content, password)
	rwm.Unlock()
	if stmt != nil {
		stmt.Close()
	}
	CheckErr(err)
	id, err := res.LastInsertId()
	CheckErr(err)
	if id > 0 {
		return true
	}
	return false
}

//数据库搜索
func selectDb(db *sql.DB, sql string) (string, string, string, string) {
	rows, err := db.Query(sql)
	defer func() {
		if rows != nil {
			rows.Close()
		}
	}()
	CheckErr(err)
	var (
		uid          int
		fileNameb256 string
		fileNameb    string
		contentb     string
		password     string
	)
	rows.Next()
	rows.Scan(&uid, &fileNameb256, &fileNameb, &contentb, &password)
	return fileNameb256, fileNameb, contentb, password
}

//删除数据项
func DelData(db *sql.DB, fileName256 string) {
	stmt3, err := db.Prepare("delete from secret_key where fileName256=?")
	defer func() {
		if stmt3 != nil {
			stmt3.Close()
		}
	}()
	CheckErr(err)
	rwm.Lock() //加锁
	_, err = stmt3.Exec(fileName256)
	rwm.Unlock() //解锁
	CheckErr(err)
}

//加密过程
func jiami(fileName string, db *sql.DB, jiamiPw string, jiajie string) {
	// delfileName := fileName
	waiB := 0
	fileSha256 := GetSha256(fileName) //文件名改为哈希256
	sql1 := fmt.Sprintf(`select * from secret_key where fileName256="%s"`, fileSha256)
	sfileSha256, _, sbeforeNegation, _ := selectDb(db, sql1) //查找数据库文件名是否已存在
	srcFile, err := os.Open(fileName)                        //读取原文件
	if err != nil {                                          //读取错误
		// fmt.Println(fileName, " -> 打开错误！")
		if waiB == 0 {
			waitgroup.Done()
			waiB++
		}
		return
	}
	dstFile, err2 := os.OpenFile(fileSha256, os.O_WRONLY|os.O_CREATE, 0666) //加密后的文件
	if err2 != nil {                                                        //读取错误
		// fmt.Println(fileSha256, " -> 打开错误！")
		if waiB == 0 {
			waitgroup.Done()
			waiB++
		}
		return
	}
	defer func() { //关闭打开的文件
		srcFile.Close()
		dstFile.Close()
	}()
	r := bufio.NewReader(srcFile) //读文件
	w := bufio.NewWriter(dstFile) //写文件
	buf := make([]byte, 4096)     //读写缓冲区
	var sBuf []byte               //加密后的数据
	bbb := 0                      //计数器，第一次读到的文件内容存储到数据库，后面内容存储到文件中
	for {
		n, err3 := r.Read(buf) //读文件到缓冲区
		if err3 != nil {       //读文件错误
			if err3 == io.EOF { //读取完毕
				w.Flush()
				fmt.Println(fileName, " -> 完成！")
				DelFile(srcFile, fileName, jiajie) //删除原文件
			} else { //非读取完毕错误
				// fmt.Println(fileName, " -> 读取错误！")

			}
			if waiB == 0 {
				waitgroup.Done()
				waiB++
			}
			return
		} else { //读取正常
			if bbb == 0 {
				//加密保存数据库
				beforeNegation := string(GetNegation(buf[:n])) // 加密数据库文件内容
				if sfileSha256 != "" {
					if sbeforeNegation == beforeNegation { //数据库有相同内容
						if IsFileExist(sfileSha256) { //加密后文件存在
							fmt.Println(fileName, "：文件已加密！")
							DelFile(srcFile, fileName, jiajie) //直接删除原文件
							if waiB == 0 {
								waitgroup.Done()
								waiB++
							}
							return
						}
					} else { //数据库没有相同内容，只是文件名相同
						ns := time.Now().UnixNano()
						fileName = fmt.Sprintf("%v%s", ns, fileName) // 重命名文件名
						fileSha256 = GetSha256(fileName)             //文件名哈希256
					}
				}
				fileNegation := string(GetNegation([]byte(fileName)))                     //原文件名加密存储到数据库
				password := GetSha256(jiamiPw)                                            //密码的哈希值
				if SaveFileToDb(db, fileSha256, fileNegation, beforeNegation, password) { //保存内容到数据库
					// fmt.Println(fileName, " -> 数据库保存成功！")
				}
			} else {
				//加密保存文件
				sBuf = GetNegation(buf[:n]) //加密数据
				_, err4 := w.Write(sBuf)    //保存数据
				if err4 != nil {            //保存错误
					fmt.Println(fileSha256, " -> 保存错误！")
					if waiB == 0 {
						waitgroup.Done()
						waiB++
					}
					return
				}
			}
		}
		bbb++ //计数器加一
	}
}

// 解密过程
func jiemi(fileName string, db *sql.DB, jiemiPw string, jiajie string) {
	sql := fmt.Sprintf(`select * from secret_key where fileName256="%s"`, fileName)
	fileNameDb256, fileNameDb, contentDb, passwordDb := selectDb(db, sql) //查找数据库
	waiB := 0
	if fileNameDb256 != "" && GetSha256(jiemiPw) == passwordDb { //数据存在加密文件，并验证密码
		fileNameDbs := string(GetNegation([]byte(fileNameDb))) //解密后文件名
		contentDbs := GetNegation([]byte(contentDb))           //解密后前半内容
		//文件名已存在，判断前半内容是否相同，若相同，直接删除加密文件、数据库内容
		if IsFileExist(fileNameDbs) { //解密后文件名已存在
			srcFile, err := os.Open(fileNameDbs) //读取原文件
			if err != nil {                      //读取错误
				// fmt.Println(fileName, " -> 打开错误！")
				if waiB == 0 {
					waitgroup.Done()
					waiB++
				}
				return
			}
			defer func() { //关闭打开的文件
				srcFile.Close()
			}()
			r := bufio.NewReader(srcFile) //读文件
			buf := make([]byte, 4096)     //读写缓冲区
			n, err3 := r.Read(buf)        //读文件到缓冲区
			if err3 != nil {              //读文件错误
				// fmt.Println(fileName, " -> 读取错误！")
				if waiB == 0 {
					waitgroup.Done()
					waiB++
				}
				return
			} else { //读取正常
				if string(buf[:n]) == string(contentDbs) {
					//解密后文件已存在
					DelData(db, fileNameDb256) //解密文件已存在，删除数据库中内容
					if waiB == 0 {
						waitgroup.Done()
						waiB++
					}
					return
				} else {
					//解密后文件名存在，但是文件内容不同
					ns := time.Now().UnixNano()
					fileNameDbs = fmt.Sprintf("%v%s", ns, fileName) // 重命名文件名
				}
			}
		}
		f, err2 := os.Open(fileName)
		if err2 != nil {
			// fmt.Println(fileName, " -> 打开错误！")
			if waiB == 0 {
				waitgroup.Done()
				waiB++
			}
			return
		}
		f2, err3 := os.OpenFile(fileNameDbs, os.O_WRONLY|os.O_CREATE, 0666)
		if err3 != nil {
			// fmt.Println(fileNameDbs, " -> 打开错误！")
			if waiB == 0 {
				waitgroup.Done()
				waiB++
			}
			return
		}
		defer func() {
			f.Close()
			f2.Close()
		}()
		r := bufio.NewReader(f)
		w := bufio.NewWriter(f2)
		buf2 := make([]byte, 4096)
		ddd := 0
		var sBuf []byte
		for {
			if ddd == 0 {
				w.Write(contentDbs)
			} else {
				n, err4 := r.Read(buf2)
				if err4 != nil {
					if err4 == io.EOF {
						w.Flush()
						fmt.Println(fileName, " -> 完成！")
						DelFile(f, fileNameDb256, jiajie) //解密成功，删除加密后的文件
						DelData(db, fileNameDb256)        //解密成功，删除数据库中内容

					} else { //非读取完毕错误
						// fmt.Println(fileName, " -> 读取错误！")
					}
					if waiB == 0 {
						waitgroup.Done()
						waiB++
					}
					return
				} else {
					sBuf = GetNegation(buf2[:n]) //解密数据
					w.Write(sBuf)
				}
			}
			ddd++
		}
	} else {
		fmt.Println(fileName, " -> 密码错误！")
	}
	if waiB == 0 {
		waitgroup.Done()
		waiB++
	}
}

var waitgroup sync.WaitGroup
var rwm sync.RWMutex

//主函数
func main() {
	db, err := sql.Open("sqlite3", "./重要勿删")
	defer db.Close()
	if err != nil {
		// fmt.Println("打开数据库错误！")
	}
	sqlTable := `
        CREATE TABLE IF NOT EXISTS secret_key(
            uid INTEGER PRIMARY KEY AUTOINCREMENT,
            fileName256 VARCHAR NULL,
            fileName VARCHAR NULL,
	        content VARCHAR NULL,
			password VARCHAR NULL
        );
	`
	db.Exec(sqlTable)
	fileDir := "."         //需要加密的相对路径
	var selectMenus string //选择菜单
	var extlist []string   //加密扩展名列表
	var ext string         //输入的加密扩展名
	var jiemiPw string
	var jiamiPw string
	for {
		selectMenus = ""
		fmt.Println("\n\n************************* 加密、解密小程序 *************************")
		fmt.Print("1. 加密    2. 解密    3. 退出\n请输入数字并回车：")
		fmt.Scanln(&selectMenus) //主菜单
		if selectMenus == "1" {  //选择加密
			extlist = []string{} //加密扩展名清空
			fmt.Print("请添加一个需要加密的扩展名（如：jpg）输入后回车：")
			for { //循环输入多个扩展名
				fmt.Scanln(&ext)                //输入需要加密的扩展名
				if ext == "11" || ext == "33" { //当等于11，33时退出输入扩展名
					break
				}
				if ext == "22" { //22清空输入的扩展名
					extlist = []string{}
					fmt.Print("请重新输入需要加密的扩展名（如：jpg）：")
					continue
				}
				extlist = append(extlist, ext) //扩展名添加到列表
				fmt.Print("您输入的扩展名为：")
				for _, i := range extlist { //循环输出输入的扩展名
					fmt.Print(i, "  ")
				}
				fmt.Print("（11.开始加密  22.重新输入 33.退出加密）：")
			}
			// extlist = []string{"jpg", "mp4"}
			if ext == "33" { //33时返回主菜单
				continue
			}
			fmt.Print("请输入加密密码：")
			fmt.Scanln(&jiamiPw)
			for _, fileName := range ReadFolder(fileDir, extlist) {
				waitgroup.Add(1)
				go jiami(fileName, db, jiamiPw, "加密")
			}
			waitgroup.Wait()

		} else if selectMenus == "2" {
			fmt.Print("请输入解密密码：")
			fmt.Scan(&jiemiPw)
			for _, fileName := range ReadFolderNeg(fileDir) { //读取已加密的文件列表
				waitgroup.Add(1)
				go jiemi(fileName, db, jiemiPw, "解密") //解密
			}
			waitgroup.Wait()
		} else if selectMenus == "3" {
			break
		} else {
			fmt.Println("输入错误，请重新选择！")
		}
	}
}