From 72b8d919fe81288de10b6e1c4f28020c7040f88c Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Tue, 10 Aug 2021 12:34:27 +0300
Subject: [PATCH] [#755] neofs-adm: allow to filter containers by ID

`--cid <cid1> --cid <cid2>` as well as `--cid <cid1>,<cid2>` is supported.

Signed-off-by: Evgenii Stratonikov <evgeniy@nspcc.ru>
---
 .../internal/modules/morph/container.go       | 48 +++++++++++++++++++
 cmd/neofs-adm/internal/modules/morph/root.go  |  3 ++
 2 files changed, 51 insertions(+)

diff --git a/cmd/neofs-adm/internal/modules/morph/container.go b/cmd/neofs-adm/internal/modules/morph/container.go
index 2d3fbec12..de063c475 100644
--- a/cmd/neofs-adm/internal/modules/morph/container.go
+++ b/cmd/neofs-adm/internal/modules/morph/container.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"io/ioutil"
+	"sort"
 
 	"github.com/nspcc-dev/neo-go/pkg/crypto/hash"
 	"github.com/nspcc-dev/neo-go/pkg/io"
@@ -67,9 +68,17 @@ func dumpContainers(cmd *cobra.Command, _ []string) error {
 		cids = append(cids, id)
 	}
 
+	isOK, err := getCIDFilterFunc(cmd)
+	if err != nil {
+		return err
+	}
+
 	var containers []*Container
 	bw := io.NewBufBinWriter()
 	for _, id := range cids {
+		if !isOK(id) {
+			continue
+		}
 		bw.Reset()
 		emit.AppCall(bw.BinWriter, ch, "get", callflag.All, id)
 		emit.AppCall(bw.BinWriter, ch, "eACL", callflag.All, id)
@@ -138,9 +147,17 @@ func restoreContainers(cmd *cobra.Command, _ []string) error {
 		return fmt.Errorf("can't parse dump file: %w", err)
 	}
 
+	isOK, err := getCIDFilterFunc(cmd)
+	if err != nil {
+		return err
+	}
+
 	bw := io.NewBufBinWriter()
 	for _, cnt := range containers {
 		hv := hash.Sha256(cnt.Value)
+		if !isOK(hv[:]) {
+			continue
+		}
 		bw.Reset()
 		emit.AppCall(bw.BinWriter, ch, "get", callflag.All, hv.BytesBE())
 		res, err := wCtx.Client.InvokeScript(bw.Bytes(), nil)
@@ -285,3 +302,34 @@ func (c *EACL) FromStackItem(item stackitem.Item) error {
 	c.Token = tok
 	return nil
 }
+
+// getCIDFilterFunc returns filtering function for container IDs.
+// Raw byte slices are used because it works with structures returned
+// from contract.
+func getCIDFilterFunc(cmd *cobra.Command) (func([]byte) bool, error) {
+	rawIDs, err := cmd.Flags().GetStringSlice(containerIDsFlag)
+	if err != nil {
+		return nil, err
+	}
+	if len(rawIDs) == 0 {
+		return func([]byte) bool { return true }, nil
+	}
+
+	for i := range rawIDs {
+		err := cid.New().Parse(rawIDs[i])
+		if err != nil {
+			return nil, fmt.Errorf("can't parse CID %s: %w", rawIDs[i], err)
+		}
+	}
+	sort.Strings(rawIDs)
+	return func(rawID []byte) bool {
+		var v [32]byte
+		copy(v[:], rawID)
+
+		id := cid.New()
+		id.SetSHA256(v)
+		idStr := id.String()
+		n := sort.Search(len(rawIDs), func(i int) bool { return rawIDs[i] >= idStr })
+		return n < len(rawIDs) && rawIDs[n] == idStr
+	}, nil
+}
diff --git a/cmd/neofs-adm/internal/modules/morph/root.go b/cmd/neofs-adm/internal/modules/morph/root.go
index 928cb470d..6f78bfb78 100644
--- a/cmd/neofs-adm/internal/modules/morph/root.go
+++ b/cmd/neofs-adm/internal/modules/morph/root.go
@@ -29,6 +29,7 @@ const (
 	withdrawFeeCLIFlag    = "withdraw-fee"
 	containerDumpFlag     = "dump"
 	containerContractFlag = "container-contract"
+	containerIDsFlag      = "cid"
 )
 
 var (
@@ -171,9 +172,11 @@ func init() {
 	dumpContainersCmd.Flags().StringP(endpointFlag, "r", "", "N3 RPC node endpoint")
 	dumpContainersCmd.Flags().String(containerDumpFlag, "", "file where to save dumped containers")
 	dumpContainersCmd.Flags().String(containerContractFlag, "", "container contract hash (for networks without NNS)")
+	dumpContainersCmd.Flags().StringSlice(containerIDsFlag, nil, "containers to dump")
 
 	RootCmd.AddCommand(restoreContainersCmd)
 	restoreContainersCmd.Flags().String(alphabetWalletsFlag, "", "path to alphabet wallets dir")
 	restoreContainersCmd.Flags().StringP(endpointFlag, "r", "", "N3 RPC node endpoint")
 	restoreContainersCmd.Flags().String(containerDumpFlag, "", "file to restore containers from")
+	restoreContainersCmd.Flags().StringSlice(containerIDsFlag, nil, "containers to restore")
 }