github.com/letsencrypt/trillian@v1.1.2-0.20180615153820-ae375a99d36a/cmd/updatetree/main.go (about)

     1  // Copyright 2018 Google Inc. All Rights Reserved.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package main contains the implementation and entry point for the updatetree
    16  // command.
    17  //
    18  // Example usage:
    19  // $ ./updatetree --admin_server=host:port --tree_id=123456789 --tree_state=FROZEN
    20  //
    21  // The output is minimal to allow for easy usage in automated scripts.
    22  package main
    23  
    24  import (
    25  	"context"
    26  	"errors"
    27  	"flag"
    28  	"fmt"
    29  	"time"
    30  
    31  	"github.com/golang/glog"
    32  	"github.com/golang/protobuf/proto"
    33  	"github.com/google/trillian"
    34  	"google.golang.org/genproto/protobuf/field_mask"
    35  	"google.golang.org/grpc"
    36  	"google.golang.org/grpc/codes"
    37  	"google.golang.org/grpc/status"
    38  )
    39  
    40  var (
    41  	adminServerAddr = flag.String("admin_server", "", "Address of the gRPC Trillian Admin Server (host:port)")
    42  	rpcDeadline     = flag.Duration("rpc_deadline", time.Second*10, "Deadline for RPC requests")
    43  	treeID          = flag.Int64("tree_id", 0, "The ID of the tree to be set updated")
    44  	treeState       = flag.String("tree_state", "", "If set the tree state will be updated")
    45  )
    46  
    47  // TODO(Martin2112): Pass everything needed into this and don't refer to flags.
    48  func updateTree(ctx context.Context) (*trillian.Tree, error) {
    49  	if *adminServerAddr == "" {
    50  		return nil, errors.New("empty --admin_server, please provide the Admin server host:port")
    51  	}
    52  
    53  	m := proto.EnumValueMap("trillian.TreeState")
    54  	if m == nil {
    55  		return nil, fmt.Errorf("can't find enum value map for states")
    56  	}
    57  	newState, ok := m[*treeState]
    58  	if !ok {
    59  		return nil, fmt.Errorf("invalid tree state: %v", *treeState)
    60  	}
    61  
    62  	// We only want to update the state of the tree, which means we need a field
    63  	// mask on the request.
    64  	treeStateMask := &field_mask.FieldMask{
    65  		Paths: []string{"tree_state"},
    66  	}
    67  
    68  	req := &trillian.UpdateTreeRequest{
    69  		Tree: &trillian.Tree{
    70  			TreeId:    *treeID,
    71  			TreeState: trillian.TreeState(newState),
    72  		},
    73  		UpdateMask: treeStateMask,
    74  	}
    75  
    76  	conn, err := grpc.Dial(*adminServerAddr, grpc.WithInsecure())
    77  	if err != nil {
    78  		return nil, fmt.Errorf("failed to dial %v: %v", *adminServerAddr, err)
    79  	}
    80  	defer conn.Close()
    81  
    82  	client := trillian.NewTrillianAdminClient(conn)
    83  	for {
    84  		tree, err := client.UpdateTree(ctx, req)
    85  		if err == nil {
    86  			return tree, nil
    87  		}
    88  		if s, ok := status.FromError(err); ok && s.Code() == codes.Unavailable {
    89  			glog.Errorf("Admin server unavailable, trying again: %v", err)
    90  			time.Sleep(100 * time.Millisecond)
    91  			continue
    92  		}
    93  		return nil, fmt.Errorf("failed to UpdateTree(%+v): %T %v", req, err, err)
    94  	}
    95  }
    96  
    97  func main() {
    98  	flag.Parse()
    99  	defer glog.Flush()
   100  
   101  	ctx, cancel := context.WithTimeout(context.Background(), *rpcDeadline)
   102  	defer cancel()
   103  	tree, err := updateTree(ctx)
   104  	if err != nil {
   105  		glog.Exitf("Failed to update tree: %v", err)
   106  	}
   107  
   108  	// DO NOT change the output format, scripts are meant to depend on it.
   109  	// If you really want to change it, provide an output_format flag and
   110  	// keep the default as-is.
   111  	fmt.Println(tree.TreeState)
   112  }