Gitblit devel
view | history | commit | commitdiff
James Moger
2015-09-18 d41034ad2ff334707a58ace5f3cd1de5fcfa1815
commit | author | age
521cb6 1 /*
JM 2  * Copyright 2014 gitblit.com.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *     http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 package com.gitblit.tests;
17
18 import java.io.ByteArrayInputStream;
19 import java.io.ByteArrayOutputStream;
20 import java.io.IOException;
21 import java.io.OutputStreamWriter;
22 import java.io.Writer;
23 import java.net.SocketAddress;
24 import java.security.KeyPair;
25 import java.security.KeyPairGenerator;
26 import java.security.PublicKey;
5485da 27 import java.util.ArrayList;
FB 28 import java.util.List;
521cb6 29 import java.util.concurrent.atomic.AtomicBoolean;
JM 30
31 import org.apache.sshd.client.ServerKeyVerifier;
d41034 32 import org.apache.sshd.client.SshClient;
JM 33 import org.apache.sshd.client.auth.UserAuth;
34 import org.apache.sshd.client.auth.UserAuthPublicKeyFactory;
35 import org.apache.sshd.client.channel.ClientChannel;
36 import org.apache.sshd.client.session.ClientSession;
5485da 37 import org.apache.sshd.common.NamedFactory;
521cb6 38 import org.apache.sshd.common.util.SecurityUtils;
JM 39 import org.junit.After;
40 import org.junit.AfterClass;
41 import org.junit.Before;
42 import org.junit.BeforeClass;
43
44 import com.gitblit.Constants.AccessPermission;
45 import com.gitblit.transport.ssh.IPublicKeyManager;
46 import com.gitblit.transport.ssh.MemoryKeyManager;
47 import com.gitblit.transport.ssh.SshKey;
48
49 /**
50  * Base class for SSH unit tests.
51  */
52 public abstract class SshUnitTest extends GitblitUnitTest {
53
54     protected static final AtomicBoolean started = new AtomicBoolean(false);
55     protected static KeyPairGenerator generator;
56     protected KeyPair rwKeyPair;
57     protected KeyPair roKeyPair;
58     protected String username = "admin";
59     protected String password = "admin";
60
61     @BeforeClass
62     public static void startGitblit() throws Exception {
63         generator = SecurityUtils.getKeyPairGenerator("RSA");
64         started.set(GitBlitSuite.startGitblit());
65     }
66
67     @AfterClass
68     public static void stopGitblit() throws Exception {
69         if (started.get()) {
70             GitBlitSuite.stopGitblit();
71         }
72     }
73
74     protected MemoryKeyManager getKeyManager() {
75         IPublicKeyManager mgr = gitblit().getPublicKeyManager();
76         if (mgr instanceof MemoryKeyManager) {
77             return (MemoryKeyManager) gitblit().getPublicKeyManager();
78         } else {
79             throw new RuntimeException("unexpected key manager type " + mgr.getClass().getName());
80         }
81     }
82
83     @Before
84     public void prepare() {
85         rwKeyPair = generator.generateKeyPair();
86
87         MemoryKeyManager keyMgr = getKeyManager();
88         keyMgr.addKey(username, new SshKey(rwKeyPair.getPublic()));
89
90         roKeyPair = generator.generateKeyPair();
91         SshKey sshKey = new SshKey(roKeyPair.getPublic());
92         sshKey.setPermission(AccessPermission.CLONE);
93         keyMgr.addKey(username, sshKey);
94     }
95
96     @After
97     public void tearDown() {
98         MemoryKeyManager keyMgr = getKeyManager();
99         keyMgr.removeAllKeys(username);
100     }
101
102     protected SshClient getClient() {
103         SshClient client = SshClient.setUpDefaultClient();
104         client.setServerKeyVerifier(new ServerKeyVerifier() {
105             @Override
106             public boolean verifyServerKey(ClientSession sshClientSession, SocketAddress remoteAddress, PublicKey serverKey) {
107                 return true;
108             }
109         });
5485da 110         List<NamedFactory<UserAuth>> userAuthFactories = new ArrayList<>();
d41034 111         userAuthFactories.add(new UserAuthPublicKeyFactory());
5485da 112         client.setUserAuthFactories(userAuthFactories);
521cb6 113         client.start();
JM 114         return client;
115     }
116
117     protected String testSshCommand(String cmd) throws IOException, InterruptedException {
118         return testSshCommand(cmd, null);
119     }
120
121     protected String testSshCommand(String cmd, String stdin) throws IOException, InterruptedException {
122         SshClient client = getClient();
123         ClientSession session = client.connect(username, "localhost", GitBlitSuite.sshPort).await().getSession();
124         session.addPublicKeyIdentity(rwKeyPair);
125         assertTrue(session.auth().await().isSuccess());
126
127         ClientChannel channel = session.createChannel(ClientChannel.CHANNEL_EXEC, cmd);
128         ByteArrayOutputStream baos = new ByteArrayOutputStream();
129         if (stdin != null) {
130             Writer w = new OutputStreamWriter(baos);
131             w.write(stdin);
132             w.close();
133         }
134         channel.setIn(new ByteArrayInputStream(baos.toByteArray()));
135
136         ByteArrayOutputStream out = new ByteArrayOutputStream();
137         ByteArrayOutputStream err = new ByteArrayOutputStream();
138         channel.setOut(out);
139         channel.setErr(err);
140         channel.open();
141
142         channel.waitFor(ClientChannel.CLOSED, 0);
143
144         String result = out.toString().trim();
145         channel.close(false);
146         client.stop();
147         return result;
148     }
149 }