-
Notifications
You must be signed in to change notification settings - Fork 0
/
pipeline_cmd.go
144 lines (128 loc) · 3.79 KB
/
pipeline_cmd.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
package cmd
import (
"encoding/hex"
"flag"
"fmt"
"io"
"log"
"os"
"strconv"
"github.com/OhanaFS/stitch"
"github.com/OhanaFS/stitch/util"
)
var (
PipelineCmd = flag.NewFlagSet("pipeline", flag.ExitOnError)
plInputFile = PipelineCmd.String("input", "", "path to the input file")
plOutputFile = PipelineCmd.String("output", "", "path to the output file")
plDataShards = PipelineCmd.Int("data-shards", 2, "number of data shards")
plParityShards = PipelineCmd.Int("parity-shards", 1, "number of parity shards")
plFileKey = PipelineCmd.String("file-key", "00000000000000000000000000000000", "file key")
plFileKeySalt = PipelineCmd.String("file-key-salt", "000000000000000000000000", "file key salt")
)
func RunPipelineCmd() int {
// Make sure the user specifies either the input or output file.
isInput := *plInputFile != ""
isOutput := *plOutputFile != ""
if isInput == isOutput {
log.Fatalln("You must specify either -input or -output.")
}
// Grab the filename
fileName := *plInputFile
if isOutput {
fileName = *plOutputFile
}
// Generate a list of shard names
totalShards := *plDataShards + *plParityShards
shardNames := make([]string, totalShards)
for i := 0; i < totalShards; i++ {
shardNames[i] = fileName + ".shard" + strconv.Itoa(i)
}
// Create the encoder
encoder := stitch.NewEncoder(&stitch.EncoderOptions{
DataShards: uint8(*plDataShards),
ParityShards: uint8(*plParityShards),
KeyThreshold: uint8(*plDataShards),
})
// Get key and IV
key, err := hex.DecodeString(*plFileKey)
if err != nil {
log.Fatalln("Invalid key:", err)
}
iv, err := hex.DecodeString(*plFileKeySalt)
if err != nil {
log.Fatalln("Invalid IV:", err)
}
if isInput {
// Open a file for reading
fmt.Printf("Opening file %s for reading\n", fileName)
file, err := os.Open(fileName)
if err != nil {
log.Fatalln("Failed to open file:", err)
}
defer file.Close()
// Open the output files
shardWriters := make([]io.Writer, totalShards)
shardFiles := make([]*os.File, totalShards)
for i := 0; i < totalShards; i++ {
shardFile, err := os.Create(shardNames[i])
if err != nil {
log.Fatalf("Failed to create shard %d: %s\n", i, err)
}
defer shardFile.Close()
shardWriters[i] = shardFile
shardFiles[i] = shardFile
}
// Set up progress bar
stat, err := file.Stat()
if err != nil {
log.Fatalln("Failed to stat file:", err)
}
progressReader := util.NewProgressReader(file, stat.Size())
// Encode the file
log.Println("Encoding file...")
if _, err = encoder.Encode(progressReader, shardWriters, key, iv); err != nil {
log.Fatalln("Failed to encode file:", err)
}
fmt.Println("")
// Finalize shard headers
log.Println("Finalizing shard headers...")
for i := 0; i < totalShards; i++ {
if err = encoder.FinalizeHeader(shardFiles[i]); err != nil {
log.Fatalf("Failed to finalize shard %d: %s\n", i, err)
}
}
} else {
// Open output file for writing
fmt.Printf("Opening file %s for writing\n", fileName)
outputFile, err := os.Create(fileName)
if err != nil {
log.Fatalln("Failed to open output file:", err)
}
defer outputFile.Close()
// Open the input files
var shards []io.ReadSeeker
for i := 0; i < totalShards; i++ {
shardFile, err := os.Open(shardNames[i])
if err != nil {
log.Println("Failed to open shard:", err)
continue
}
defer shardFile.Close()
shards = append(shards, shardFile)
}
// Decode the file
log.Println("Decoding file...")
reader, err := encoder.NewReadSeeker(shards, key, iv)
if err != nil {
log.Fatalln("Failed to create reader:", err)
}
n, err := io.Copy(outputFile, util.NewProgressReader(reader, 0))
if err != nil {
log.Fatalln("Failed to decode file:", err)
}
fmt.Println("")
log.Printf("Decoded %d bytes\n", n)
}
log.Println("Done.")
return 0
}