From 32af7251f2aeb04267f3c5b4c3be62647aa3ece2 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Thu, 24 Feb 2022 15:30:47 +0300
Subject: [PATCH] [#134] owner: Add helpers for working script hash

Signed-off-by: Evgenii Stratonikov <evgeniy@nspcc.ru>
---
 owner/convert_test.go | 12 ++++++++++++
 owner/id.go           | 11 +++++++++++
 owner/id_test.go      |  9 +++++++++
 3 files changed, 32 insertions(+)

diff --git a/owner/convert_test.go b/owner/convert_test.go
index 6c20507..494285f 100644
--- a/owner/convert_test.go
+++ b/owner/convert_test.go
@@ -32,3 +32,15 @@ func TestPublicKeyToBytes(t *testing.T) {
 	require.Equal(t, expected, actual)
 	require.Equal(t, NEO3WalletSize, len(actual))
 }
+
+func TestScriptHashToIDBytes(t *testing.T) {
+	p, err := keys.NewPrivateKey()
+	require.NoError(t, err)
+
+	expected, err := base58.Decode(p.PublicKey().Address())
+	require.NoError(t, err)
+
+	actual := ScriptHashToIDBytes(p.GetScriptHash())
+	require.Equal(t, expected, actual)
+	require.Equal(t, NEO3WalletSize, len(actual))
+}
diff --git a/owner/id.go b/owner/id.go
index bdc355b..98c2de9 100644
--- a/owner/id.go
+++ b/owner/id.go
@@ -10,6 +10,7 @@ import (
 	"github.com/nspcc-dev/neo-go/pkg/crypto/hash"
 	"github.com/nspcc-dev/neo-go/pkg/crypto/keys"
 	"github.com/nspcc-dev/neo-go/pkg/encoding/address"
+	"github.com/nspcc-dev/neo-go/pkg/util"
 	"github.com/nspcc-dev/neo-go/pkg/wallet"
 	"github.com/nspcc-dev/neofs-api-go/v2/refs"
 )
@@ -47,6 +48,11 @@ func (id *ID) SetPublicKey(pub *ecdsa.PublicKey) {
 	(*refs.OwnerID)(id).SetValue(PublicKeyToIDBytes(pub))
 }
 
+// SetScriptHash sets owner identifier value to the provided NEO3 script hash.
+func (id *ID) SetScriptHash(u util.Uint160) {
+	(*refs.OwnerID)(id).SetValue(ScriptHashToIDBytes(u))
+}
+
 // ToV2 returns the v2 owner ID message.
 //
 // Nil ID converts to nil.
@@ -144,6 +150,11 @@ func (id *ID) UnmarshalJSON(data []byte) error {
 // make it clear that no errors can occur.
 func PublicKeyToIDBytes(pub *ecdsa.PublicKey) []byte {
 	sh := (*keys.PublicKey)(pub).GetScriptHash()
+	return ScriptHashToIDBytes(sh)
+}
+
+// ScriptHashToIDBytes converts NEO3 script hash to a byte slice of NEO3WalletSize length.
+func ScriptHashToIDBytes(sh util.Uint160) []byte {
 	b := make([]byte, NEO3WalletSize)
 	b[0] = address.Prefix
 	copy(b[1:], sh.BytesBE())
diff --git a/owner/id_test.go b/owner/id_test.go
index c77d200..a27aec8 100644
--- a/owner/id_test.go
+++ b/owner/id_test.go
@@ -163,3 +163,12 @@ func TestNewID(t *testing.T) {
 		require.Nil(t, idV2.GetValue())
 	})
 }
+
+func TestID_SetScriptHash(t *testing.T) {
+	p, err := keys.NewPrivateKey()
+	require.NoError(t, err)
+
+	id := NewID()
+	id.SetScriptHash(p.GetScriptHash())
+	require.Equal(t, p.Address(), id.String())
+}