PY-14605 Support remote single-user ipnb servers

TODO: Get ipnb file path correctly as a relative path from working directory

python/ipnb/src/org/jetbrains/plugins/ipnb/configuration/IpnbConnectionManager.java

@@ -51,6 +51,8 @@ import org.jetbrains.plugins.ipnb.protocol.IpnbConnectionV3;
 import javax.swing.event.HyperlinkEvent;
 import java.io.IOException;
 import java.net.*;
+import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -114,17 +116,21 @@ public final class IpnbConnectionManager implements ProjectComponent {
                                          @NotNull final IpnbFileEditor fileEditor,
                                          @NotNull final String path,
                                          @NotNull final String url) {
-    if (myToken != null) return startConnection(codePanel, path, url, false);
-    final Module module = ProjectFileIndex.SERVICE.getInstance(myProject).getModuleForFile(fileEditor.getVirtualFile());
-    if (module != null) {
-      final Sdk sdk = PythonSdkType.findPythonSdk(module);
-      if (sdk != null) {
-        final List<PyPackage> packages = PyPackageManager.getInstance(sdk).getPackages();
-        if (packages != null) {
-          final PyPackage notebookPackage = PyPackageUtil.findPackage(packages, "notebook");
-          if (notebookPackage != null && VersionComparatorUtil.compare(notebookPackage.getVersion(), "4.3.0") >= 0) {
-            myToken = askForToken(url);
-            if (myToken == null) return false;
+    final IpnbSettings ipnbSettings = IpnbSettings.getInstance(myProject);
+    final boolean isRemote = !ipnbSettings.getUsername().isEmpty() && !ipnbSettings.getPassword().isEmpty();
+    if (!isRemote) {
+      if (myToken != null) return startConnection(codePanel, path, url, false);
+      final Module module = ProjectFileIndex.SERVICE.getInstance(myProject).getModuleForFile(fileEditor.getVirtualFile());
+      if (module != null) {
+        final Sdk sdk = PythonSdkType.findPythonSdk(module);
+        if (sdk != null) {
+          final List<PyPackage> packages = PyPackageManager.getInstance(sdk).getPackages();
+          if (packages != null) {
+            final PyPackage notebookPackage = PyPackageUtil.findPackage(packages, "notebook");
+            if (notebookPackage != null && VersionComparatorUtil.compare(notebookPackage.getVersion(), "4.3.0") >= 0) {
+              myToken = askForToken(url);
+              if (myToken == null) return false;
+            }
           }
         }
       }
@@ -271,10 +277,30 @@ public final class IpnbConnectionManager implements ProjectComponent {
   private IpnbConnection getConnection(@Nullable final IpnbCodePanel codePanel, @NotNull final String urlString,
                                        @NotNull final IpnbConnectionListenerBase listener)
     throws IOException, URISyntaxException {
-    if (codePanel != null && !IpnbParser.isIpythonNewFormat(codePanel.getFileEditor().getVirtualFile())) {
-      return new IpnbConnection(urlString, listener, myToken);
+    if (codePanel != null) {
+      final VirtualFile file = codePanel.getFileEditor().getVirtualFile();
+      String pathToFile = getRelativePathToFile(file);
+      if (pathToFile != null) {
+        if (!IpnbParser.isIpythonNewFormat(file)) {
+          return new IpnbConnection(urlString, listener, myToken, myProject, pathToFile);
+        }
+        return new IpnbConnectionV3(urlString, listener, myToken, myProject, pathToFile);
+      }
     }
-    return new IpnbConnectionV3(urlString, listener, myToken);
+    
+    throw new IOException("Code panel is null");
+  }
+
+  @Nullable
+  private String getRelativePathToFile(VirtualFile file) {
+    final String workingDirectory = IpnbSettings.getInstance(myProject).getWorkingDirectory();
+    final String realWorkingDir = workingDirectory.isEmpty() ? myProject.getBasePath() : workingDirectory;
+    if (realWorkingDir != null) {
+      final Path basePath = Paths.get(realWorkingDir);
+      final Path filePath = Paths.get(file.getPath());
+      return basePath.relativize(filePath).toString();
+    }
+    return null;
   }
 
   public void interruptKernel(@NotNull final String filePath) {
@@ -310,125 +336,131 @@ public final class IpnbConnectionManager implements ProjectComponent {
   }
 
   public boolean startIpythonServer(@NotNull final String initUrl, @NotNull final IpnbFileEditor fileEditor) {
-    final Module module = ProjectFileIndex.SERVICE.getInstance(myProject).getModuleForFile(fileEditor.getVirtualFile());
-    if (module == null) return false;
-    final Sdk sdk = PythonSdkType.findPythonSdk(module);
-    if (sdk == null) {
-      showWarning(fileEditor, "Please check Python Interpreter in Settings->Python Interpreter", null);
-      return false;
-    }
-    final List<PyPackage> packages = PyPackageManager.getInstance(sdk).getPackages();
-    final PyPackage ipythonPackage = packages != null ? PyPackageUtil.findPackage(packages, "ipython") : null;
-    final PyPackage jupyterPackage = packages != null ? PyPackageUtil.findPackage(packages, "jupyter") : null;
-    if (ipythonPackage == null && jupyterPackage == null) {
-      showWarning(fileEditor, "Add Jupyter to the interpreter of the current project.", null);
-      return false;
-    }
+    final String username = IpnbSettings.getInstance(myProject).getUsername();
+    final String password = IpnbSettings.getInstance(myProject).getPassword();
+    final boolean isRemote = !username.isEmpty() && !password.isEmpty();
+    if (!isRemote) {
+      final Module module = ProjectFileIndex.SERVICE.getInstance(myProject).getModuleForFile(fileEditor.getVirtualFile());
+      if (module == null) return false;
+      final Sdk sdk = PythonSdkType.findPythonSdk(module);
+      if (sdk == null) {
+        showWarning(fileEditor, "Please check Python Interpreter in Settings->Python Interpreter", null);
+        return false;
+      }
+      final List<PyPackage> packages = PyPackageManager.getInstance(sdk).getPackages();
+      final PyPackage ipythonPackage = packages != null ? PyPackageUtil.findPackage(packages, "ipython") : null;
+      final PyPackage jupyterPackage = packages != null ? PyPackageUtil.findPackage(packages, "jupyter") : null;
+      if (ipythonPackage == null && jupyterPackage == null) {
+        showWarning(fileEditor, "Add Jupyter to the interpreter of the current project.", null);
+        return false;
+      }
 
-    String url = showDialogUrl(initUrl);
-    if (url == null) return false;
-    final IpnbSettings ipnbSettings = IpnbSettings.getInstance(myProject);
-    ipnbSettings.setURL(url);
+      String url = showDialogUrl(initUrl);
+      if (url == null) return false;
+      final IpnbSettings ipnbSettings = IpnbSettings.getInstance(myProject);
+      ipnbSettings.setURL(url);
 
-    final Pair<String, String> hostPort = getHostPortFromUrl(url);
-    if (hostPort == null) {
-      showWarning(fileEditor, "Please, check Jupyter Notebook URL in <a href=\"\">Settings->Tools->Jupyter Notebook</a>",
-                  new IpnbSettingsAdapter());
-      return false;
-    }
-    final String homePath = sdk.getHomePath();
-    if (homePath == null) {
-      showWarning(fileEditor, "Python Sdk is invalid, please check Python Interpreter in Settings->Python Interpreter", null);
-      return false;
-    }
-    Map<String, String> env = null;
-    final ArrayList<String> parameters = Lists.newArrayList(homePath);
-    String ipython = findJupyterRunner(homePath);
-    if (ipython == null) {
-      ipython = findIPythonRunner(homePath);
+      final Pair<String, String> hostPort = getHostPortFromUrl(url);
+      if (hostPort == null) {
+        showWarning(fileEditor, "Please, check Jupyter Notebook URL in <a href=\"\">Settings->Tools->Jupyter Notebook</a>",
+                    new IpnbSettingsAdapter());
+        return false;
+      }
+      final String homePath = sdk.getHomePath();
+      if (homePath == null) {
+        showWarning(fileEditor, "Python Sdk is invalid, please check Python Interpreter in Settings->Python Interpreter", null);
+        return false;
+      }
+      Map<String, String> env = null;
+      final ArrayList<String> parameters = Lists.newArrayList(homePath);
+      String ipython = findJupyterRunner(homePath);
       if (ipython == null) {
-        ipython = PythonHelper.LOAD_ENTRY_POINT.asParamString();
-        env = ImmutableMap.of("PYCHARM_EP_DIST", "ipython", "PYCHARM_EP_NAME", "ipython");
+        ipython = findIPythonRunner(homePath);
+        if (ipython == null) {
+          ipython = PythonHelper.LOAD_ENTRY_POINT.asParamString();
+          env = ImmutableMap.of("PYCHARM_EP_DIST", "ipython", "PYCHARM_EP_NAME", "ipython");
+        }
+        parameters.add(ipython);
+        parameters.add("notebook");
       }
-      parameters.add(ipython);
-      parameters.add("notebook");
-    }
-    else {
-      parameters.add(ipython);
-    }
-    parameters.add("--no-browser");
+      else {
+        parameters.add(ipython);
+      }
+      parameters.add("--no-browser");
 
-    if (hostPort.getFirst() != null) {
-      parameters.add("--ip");
-      parameters.add(hostPort.getFirst());
-    }
-    if (hostPort.getSecond() != null) {
-      parameters.add("--port");
-      parameters.add(hostPort.getSecond());
-    }
-    final String arguments = ipnbSettings.getArguments();
-    if (!StringUtil.isEmptyOrSpaces(arguments)) {
-      parameters.addAll(StringUtil.split(arguments, " "));
-    }
+      if (hostPort.getFirst() != null) {
+        parameters.add("--ip");
+        parameters.add(hostPort.getFirst());
+      }
+      if (hostPort.getSecond() != null) {
+        parameters.add("--port");
+        parameters.add(hostPort.getSecond());
+      }
+      final String arguments = ipnbSettings.getArguments();
+      if (!StringUtil.isEmptyOrSpaces(arguments)) {
+        parameters.addAll(StringUtil.split(arguments, " "));
+      }
 
-    final String directory = ipnbSettings.getWorkingDirectory();
-    final String baseDir = !StringUtil.isEmptyOrSpaces(directory) ? directory :
-                           ModuleRootManager.getInstance(module).getContentRoots()[0].getCanonicalPath();
-    final GeneralCommandLine commandLine = new GeneralCommandLine(parameters).withWorkDirectory(baseDir);
-    if (env != null) {
-      commandLine.withEnvironment(env);
-    }
+      final String directory = ipnbSettings.getWorkingDirectory();
+      final String baseDir = !StringUtil.isEmptyOrSpaces(directory) ? directory :
+                             ModuleRootManager.getInstance(module).getContentRoots()[0].getCanonicalPath();
+      final GeneralCommandLine commandLine = new GeneralCommandLine(parameters).withWorkDirectory(baseDir);
+      if (env != null) {
+        commandLine.withEnvironment(env);
+      }
 
-    try {
-      final boolean[] serverStarted = {false};
-      final KillableColoredProcessHandler processHandler = new KillableColoredProcessHandler(commandLine) {
-        @Override
-        protected void doDestroyProcess() {
-          super.doDestroyProcess();
-          myKernels.clear();
-          myToken = null;
-          UnixProcessManager.sendSigIntToProcessTree(getProcess());
-        }
+      try {
+        final boolean[] serverStarted = {false};
+        final KillableColoredProcessHandler processHandler = new KillableColoredProcessHandler(commandLine) {
+          @Override
+          protected void doDestroyProcess() {
+            super.doDestroyProcess();
+            myKernels.clear();
+            myToken = null;
+            UnixProcessManager.sendSigIntToProcessTree(getProcess());
+          }
 
-        @Override
-        public void coloredTextAvailable(@NotNull @NonNls String text, @NotNull Key attributes) {
-          super.coloredTextAvailable(text, attributes);
-          if (text.toLowerCase().contains("active kernels")) {
-            serverStarted[0] = true;
+          @Override
+          public void coloredTextAvailable(@NotNull @NonNls String text, @NotNull Key attributes) {
+            super.coloredTextAvailable(text, attributes);
+            if (text.toLowerCase().contains("active kernels")) {
+              serverStarted[0] = true;
+            }
+            final String token = "?token=";
+            if (text.toLowerCase().contains(token)) {
+              myToken = text.substring(text.indexOf(token) + token.length()).trim();
+            }
           }
-          final String token = "?token=";
-          if (text.toLowerCase().contains(token)) {
-            myToken = text.substring(text.indexOf(token) + token.length()).trim();
+
+          @Override
+          public boolean isSilentlyDestroyOnClose() {
+            return true;
           }
-        }
-
-        @Override
-        public boolean isSilentlyDestroyOnClose() {
-          return true;
+        };
+        processHandler.setShouldDestroyProcessRecursively(true);
+        GuiUtils.invokeLaterIfNeeded(() -> new RunContentExecutor(myProject, processHandler)
+          .withTitle("Jupyter Notebook")
+          .withStop(() -> {
+            myKernels.clear();
+            processHandler.destroyProcess();
+            UnixProcessManager.sendSigIntToProcessTree(processHandler.getProcess());
+          }, () -> !processHandler.isProcessTerminated())
+          .withRerun(() -> startIpythonServer(url, fileEditor))
+          .withHelpId("reference.manage.py")
+          .withFilter(new UrlFilter())
+          .run(), ModalityState.defaultModalityState());
+        int countAttempt = 0;
+        while (!serverStarted[0] && countAttempt < MAX_ATTEMPTS) {
+          countAttempt += 1;
+          TimeoutUtil.sleep(1000);
         }
-      };
-      processHandler.setShouldDestroyProcessRecursively(true);
-      GuiUtils.invokeLaterIfNeeded(() -> new RunContentExecutor(myProject, processHandler)
-        .withTitle("Jupyter Notebook")
-        .withStop(() -> {
-          myKernels.clear();
-          processHandler.destroyProcess();
-          UnixProcessManager.sendSigIntToProcessTree(processHandler.getProcess());
-        }, () -> !processHandler.isProcessTerminated())
-        .withRerun(() -> startIpythonServer(url, fileEditor))
-        .withHelpId("reference.manage.py")
-        .withFilter(new UrlFilter())
-        .run(), ModalityState.defaultModalityState());
-      int countAttempt = 0;
-      while (!serverStarted[0] && countAttempt < MAX_ATTEMPTS) {
-        countAttempt += 1;
-        TimeoutUtil.sleep(1000);
+        return true;
+      }
+      catch (ExecutionException e) {
+        return false;
       }
-      return true;
-    }
-    catch (ExecutionException e) {
-      return false;
     }
+    return true;
   }
 
   @Nullable

python/ipnb/src/org/jetbrains/plugins/ipnb/protocol/IpnbConnection.java

@@ -2,7 +2,12 @@ package org.jetbrains.plugins.ipnb.protocol;
 
 import com.google.common.collect.Lists;
 import com.google.gson.*;
+import com.intellij.openapi.diagnostic.Logger;
+import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.text.StringUtil;
+import com.intellij.util.net.HTTPMethod;
+import com.intellij.util.net.ssl.CertificateManager;
+import com.jetbrains.python.psi.PyUtil;
 import org.java_websocket.client.WebSocketClient;
 import org.java_websocket.drafts.Draft;
 import org.java_websocket.drafts.Draft_17;
@@ -10,28 +15,36 @@ import org.java_websocket.handshake.ClientHandshakeBuilder;
 import org.java_websocket.handshake.ServerHandshake;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.plugins.ipnb.configuration.IpnbSettings;
 import org.jetbrains.plugins.ipnb.format.cells.output.*;
 
-import java.io.BufferedReader;
-import java.io.IOException;
-import java.io.InputStreamReader;
+import javax.net.ssl.HostnameVerifier;
+import javax.net.ssl.HttpsURLConnection;
+import javax.net.ssl.SSLContext;
+import javax.net.ssl.SSLSession;
+import java.io.*;
 import java.net.*;
-import java.util.*;
+import java.nio.charset.StandardCharsets;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.stream.Collectors;
 
 /**
  * @author vlan
- *
- * To be removed
-*/
+ *         <p>
+ *         To be removed
+ */
 public class IpnbConnection {
+  private static final Logger LOG = Logger.getInstance(IpnbConnection.class);
   protected static final String API_URL = "/api";
   protected static final String KERNELS_URL = API_URL + "/kernels";
-  protected static final String HTTP_POST = "POST";
   // TODO: Serialize cookies for the authentication message
   protected static final String authMessage = "{\"header\":{\"msg_id\":\"\", \"msg_type\":\"connect_request\"}, \"parent_header\":\"\", \"metadata\":{}," +
                                               "\"channel\":\"shell\" }";
-  public static final String HTTP_DELETE = "DELETE";
   public static final String AUTHENTICATION_NEEDED = "Authentication needed";
+  private static final String LOGIN_URL = "/login";
 
   @NotNull protected final URI myURI;
   @NotNull protected final String myKernelId;
@@ -50,20 +63,140 @@ public class IpnbConnection {
   private IpnbOutputCell myOutput;
   private int myExecCount;
   private String myXsrf;
+  private HashMap<String, String> myHeaders = new HashMap<>();
 
 
   public IpnbConnection(@NotNull String uri, @NotNull IpnbConnectionListener listener,
-                        @Nullable final String token) throws IOException, URISyntaxException {
+                        @Nullable final String token, @NotNull Project project, @NotNull String pathToFile) throws IOException, URISyntaxException {
     myURI = new URI(uri);
     myListener = listener;
     myToken = token;
     mySessionId = UUID.randomUUID().toString();
-    initXSRF();
-    myKernelId = startKernel();
+    HttpsURLConnection.setDefaultSSLSocketFactory(CertificateManager.getInstance().getSslContext().getSocketFactory());
+    HttpsURLConnection.setDefaultHostnameVerifier(new HostnameVerifier() {
+      @Override
+      public boolean verify(String s, SSLSession session) {
+        return myURI.getHost().equals(s);
+      }
+    });
+    final boolean needed = isLoginNeeded();
+    if (needed) {
+      final IpnbSettings ipnbSettings = IpnbSettings.getInstance(project);
+      final List<HttpCookie> cookies = login(ipnbSettings.getUsername(), ipnbSettings.getPassword());
+      final String cookieString =
+        cookies.stream().map(cookie -> cookie.getName() + "=" + cookie.getValue()).collect(Collectors.joining(";"));
+      myHeaders.put("Cookie", cookieString);
+      myKernelId = getExistingKernelForSession(pathToFile);
+    }
+    else {
+      initXSRF();
+      if (myToken != null) {
+        myHeaders.put("Authorization", "token " + myToken);
+      }
+      myKernelId = startKernel();
+    }
 
     initializeClients();
   }
 
+  private List<HttpCookie> login(@NotNull String username, @NotNull String password) throws IOException {
+    CookieManager cookieManager = new CookieManager();
+    CookieHandler.setDefault(cookieManager);
+    String urlParameters = "username=" + username + "&" + "password=" + password;
+    byte[] postData = urlParameters.getBytes(StandardCharsets.UTF_8);
+    final HttpsURLConnection connection = (HttpsURLConnection)new URL(myURI + LOGIN_URL).openConnection();
+    connection.setUseCaches(false);
+    connection.setRequestMethod(HTTPMethod.POST.name());
+    connection.setRequestProperty("Content-Type", "application/x-www-form-urlencoded");
+    connection.setRequestProperty("Content-Length", Integer.toString(postData.length));
+    connection.setDoOutput(true);
+
+    final OutputStream outputStream = connection.getOutputStream();
+    try (DataOutputStream wr = new DataOutputStream(outputStream)) {
+      wr.write(postData);
+      wr.flush();
+    }
+    connection.connect();
+
+    final int code = connection.getResponseCode();
+    if (code == HttpURLConnection.HTTP_OK) {
+      return cookieManager.getCookieStore().getCookies();
+    }
+    else {
+      throw new IOException("Unable to login: " + code + " " + connection.getResponseMessage());
+    }
+  }
+
+  private String getDefaultKernelName() {
+    try {
+      final String response = httpRequest(myURI + "/api/kernelspecs", HTTPMethod.GET.name());
+      final JsonObject kernelSpecs = PyUtil.as(new JsonParser().parse(response), JsonObject.class);
+      if (kernelSpecs != null && kernelSpecs.has("default")) {
+        return kernelSpecs.get("default").getAsString();
+      }
+      else {
+        LOG.warn("Got wrong kernel specs: " + response);
+      }
+    }
+    catch (IOException e) {
+      LOG.warn(e.getMessage());
+    }
+
+    return "";
+  }
+
+  private String getExistingKernelForSession(@NotNull String pathToFile) {
+    try {
+      final URLConnection connection = new URL(myURI + "/api/sessions").openConnection();
+      if (connection instanceof HttpsURLConnection) {
+        final String kernelName = getDefaultKernelName();
+        if (kernelName.isEmpty()) return "";
+        final SessionWrapper sessionWrapper = new SessionWrapper(kernelName, pathToFile);
+        final Gson gsonBuilder = new GsonBuilder().serializeNulls().create();
+        final byte[] postData = gsonBuilder.toJson(sessionWrapper).getBytes(StandardCharsets.UTF_8);
+        final HttpsURLConnection httpsConnection = (HttpsURLConnection)configureConnection((HttpURLConnection)connection, 
+                                                                                           HTTPMethod.POST.name());
+        httpsConnection.setRequestProperty("Content-Type", "application/json");
+        httpsConnection.setRequestProperty("Content-Length", Integer.toString(postData.length));
+        httpsConnection.setRequestProperty("Referer", myURI + "/notebooks/" + pathToFile);
+        httpsConnection.setUseCaches(false);
+        httpsConnection.setDoOutput(true);
+
+        final OutputStream outputStream = connection.getOutputStream();
+        try (DataOutputStream wr = new DataOutputStream(outputStream)) {
+          wr.write(postData);
+          wr.flush();
+        }
+        connection.connect();
+        final String response = getResponse(httpsConnection);
+        final SessionWrapper wrapper = gsonBuilder.fromJson(response, SessionWrapper.class);
+
+        return wrapper.kernel.id;
+      }
+      else {
+        throw new UnsupportedOperationException("Only HTTP URLs are supported");
+      }
+    }
+    catch (IOException e) {
+      LOG.warn(e);
+    }
+    return "";
+  }
+
+  private boolean isLoginNeeded() throws IOException {
+    final HttpsURLConnection connection = PyUtil.as(new URL(myURI.toString() + "/tree").openConnection(), HttpsURLConnection.class);
+    if (connection != null) {
+      connection.setInstanceFollowRedirects(false);
+      connection.connect();
+      if (connection.getResponseCode() == HttpURLConnection.HTTP_MOVED_TEMP) {
+        final String location = connection.getHeaderField("Location");
+        connection.disconnect();
+        if (location != null && location.startsWith("/login")) return true;
+      }
+    }
+    return false;
+  }
+
   private void initXSRF() {
     try {
       CookieManager cookieManager = new CookieManager();
@@ -84,7 +217,7 @@ public class IpnbConnection {
   protected void initializeClients() throws URISyntaxException {
     final Draft draft = new Draft17WithOrigin();
 
-    myShellClient = new WebSocketClient(getShellURI(), draft) {
+    myShellClient = new WebSocketClient(getShellURI(), draft, myHeaders, 0) {
       @Override
       public void onOpen(@NotNull ServerHandshake handshakeData) {
         send(authMessage);
@@ -148,22 +281,22 @@ public class IpnbConnection {
 
   @NotNull
   private String startKernel() throws IOException {
-    final String s = httpRequest(myURI + KERNELS_URL, HTTP_POST);
+    final String s = httpRequest(myURI + KERNELS_URL, HTTPMethod.POST.name());
     final Gson gson = new Gson();
     final Kernel kernel = gson.fromJson(s, Kernel.class);
     return kernel.getId();
   }
 
   protected void shutdownKernel() throws IOException {
-    httpRequest(myURI + KERNELS_URL + "/" + myKernelId, HTTP_DELETE);
+    httpRequest(myURI + KERNELS_URL + "/" + myKernelId, HTTPMethod.DELETE.name());
   }
 
   public void interrupt() throws IOException {
-    httpRequest(myURI + KERNELS_URL + "/" + myKernelId + "/interrupt", HTTP_POST);
+    httpRequest(myURI + KERNELS_URL + "/" + myKernelId + "/interrupt", HTTPMethod.POST.name());
   }
 
   public void reload() throws IOException {
-    httpRequest(myURI + KERNELS_URL + "/" + myKernelId + "/restart", HTTP_POST);
+    httpRequest(myURI + KERNELS_URL + "/" + myKernelId + "/restart", HTTPMethod.POST.name());
   }
 
   @NotNull
@@ -178,44 +311,62 @@ public class IpnbConnection {
 
   @NotNull
   protected String getWebSocketURIBase() {
-    return "ws://" + myURI.getAuthority() + KERNELS_URL + "/" + myKernelId;
+    final String scheme = myURI.getScheme();
+    String prefix = scheme.equals("http") ? "ws://" : "wss://";
+    return prefix + myURI.getAuthority() + KERNELS_URL + "/" + myKernelId;
   }
 
   @NotNull
   private String httpRequest(@NotNull String url, @NotNull String method) throws IOException {
     final URLConnection urlConnection = new URL(url).openConnection();
     if (urlConnection instanceof HttpURLConnection) {
-      final HttpURLConnection connection = (HttpURLConnection)urlConnection;
-      connection.setRequestMethod(method);
-      connection.setReadTimeout(60000);
-      if (!StringUtil.isEmptyOrSpaces(myToken)) {
-        connection.setRequestProperty("Authorization", "token " + myToken);
-      }
-      else if (!StringUtil.isEmptyOrSpaces(myXsrf)) {
-        connection.setRequestProperty("X-XSRFToken", myXsrf);
-      }
-      if (connection.getResponseCode() == HttpURLConnection.HTTP_FORBIDDEN) {
+      final HttpURLConnection connection = configureConnection((HttpURLConnection)urlConnection, method);
+      final int code = connection.getResponseCode();
+      if (code == HttpURLConnection.HTTP_FORBIDDEN) {
         throw new IOException(AUTHENTICATION_NEEDED);
       }
-      final BufferedReader reader = new BufferedReader(new InputStreamReader(connection.getInputStream(), "utf-8"));
-      try {
-        final StringBuilder builder = new StringBuilder();
-        char[] buffer = new char[4096];
-        int n;
-        while ((n = reader.read(buffer)) != -1) {
-          builder.append(buffer, 0, n);
-        }
-        return builder.toString();
-      }
-      finally {
-        reader.close();
-      }
+      return getResponse(connection);
     }
     else {
       throw new UnsupportedOperationException("Only HTTP URLs are supported");
     }
   }
 
+  @NotNull
+  private static String getResponse(HttpURLConnection connection) throws IOException {
+    final BufferedReader reader = new BufferedReader(new InputStreamReader(connection.getInputStream(), "utf-8"));
+    try {
+      final StringBuilder builder = new StringBuilder();
+      char[] buffer = new char[4096];
+      int n;
+      while ((n = reader.read(buffer)) != -1) {
+        builder.append(buffer, 0, n);
+      }
+      return builder.toString();
+    }
+    finally {
+      reader.close();
+    }
+  }
+
+  @NotNull
+  private HttpURLConnection configureConnection(HttpURLConnection urlConnection, @NotNull String method) throws ProtocolException {
+    urlConnection.setRequestMethod(method);
+    urlConnection.setReadTimeout(60000);
+    if (!StringUtil.isEmptyOrSpaces(myToken)) {
+      urlConnection.setRequestProperty("Authorization", "token " + myToken);
+    }
+    else if (!StringUtil.isEmptyOrSpaces(myXsrf)) {
+      urlConnection.setRequestProperty("X-XSRFToken", myXsrf);
+    }
+    else if (!myHeaders.isEmpty()) {
+      for (Map.Entry<String, String> entry : myHeaders.entrySet()) {
+        urlConnection.setRequestProperty(entry.getKey(), entry.getValue());
+      }
+    }
+    return urlConnection;
+  }
+
   @NotNull
   public Message createExecuteRequest(String code, String messageId) {
     final JsonObject content = new JsonObject();
@@ -471,8 +622,20 @@ public class IpnbConnection {
 
   protected class IpnbWebSocketClient extends WebSocketClient {
     protected IpnbWebSocketClient(@NotNull final URI serverUri, @NotNull final Draft draft) {
-      super(serverUri, draft, myToken != null ? Collections.singletonMap("Authorization", "token " + myToken) :
-                                                Collections.emptyMap(), 5000);
+      super(serverUri, draft, myHeaders, 10000);
+      configureSsl(serverUri);
+    }
+
+    private void configureSsl(@NotNull URI serverUri) {
+      if (serverUri.getScheme().equals("wss")) {
+        final SSLContext sslContext = CertificateManager.getInstance().getSslContext();
+        try {
+          this.setSocket(sslContext.getSocketFactory().createSocket());
+        }
+        catch (IOException e) {
+          LOG.warn(e.getMessage());
+        }
+      }
     }
 
     @Override
@@ -529,12 +692,12 @@ public class IpnbConnection {
 
     @Override
     public void onClose(int code, String reason, boolean remote) {
-
+      LOG.info("IPNB WebSocket was closed:  code " + code + " reason: " + reason);
     }
 
     @Override
     public void onError(Exception ex) {
-
+      LOG.error(ex);
     }
   }
 
@@ -545,4 +708,31 @@ public class IpnbConnection {
   public int getExecCount() {
     return myExecCount;
   }
+
+  private static class SessionWrapper {
+    NotebookWrapper notebook;
+    KernelWrapper kernel;
+
+    public SessionWrapper(String interpreterName, String filePath) {
+      kernel = new KernelWrapper(interpreterName);
+      notebook = new NotebookWrapper(filePath);
+    }
+  }
+
+  private static class KernelWrapper {
+    String id;
+    String name;
+
+    public KernelWrapper(String name) {
+      this.name = name;
+    }
+  }
+
+  private static class NotebookWrapper {
+    String path;
+
+    public NotebookWrapper(String path) {
+      this.path = path;
+    }
+  }
 }

python/ipnb/src/org/jetbrains/plugins/ipnb/protocol/IpnbConnectionV3.java

@@ -1,6 +1,7 @@
 package org.jetbrains.plugins.ipnb.protocol;
 
 import com.google.gson.Gson;
+import com.intellij.openapi.project.Project;
 import org.java_websocket.client.WebSocketClient;
 import org.java_websocket.drafts.Draft;
 import org.jetbrains.annotations.NotNull;
@@ -15,8 +16,11 @@ public class IpnbConnectionV3 extends IpnbConnection {
   private WebSocketClient myChannelsClient;
   private Thread myChannelsThread;
 
-  public IpnbConnectionV3(@NotNull String uri, @NotNull IpnbConnectionListener listener, @Nullable final String token) throws IOException, URISyntaxException {
-    super(uri, listener, token);
+  public IpnbConnectionV3(@NotNull String uri,
+                          @NotNull IpnbConnectionListener listener,
+                          @Nullable final String token,
+                          @NotNull Project project, String pathToFile) throws IOException, URISyntaxException {
+    super(uri, listener, token, project, pathToFile);
   }
 
   @Override

python/ipnb/testSrc/WebSocketConnectionTest.java

@@ -1,3 +1,4 @@
+import com.intellij.openapi.project.DefaultProjectFactory;
 import com.intellij.openapi.util.Ref;
 import junit.framework.TestCase;
 import org.jetbrains.annotations.NotNull;
@@ -11,12 +12,11 @@ import java.net.URISyntaxException;
 import java.util.List;
 
 /**
- *
  * * Message Spec
- *   http://ipython.org/ipython-doc/dev/development/messaging.html
- *
+ * http://ipython.org/ipython-doc/dev/development/messaging.html
+ * <p>
  * * Notebook REST API
- *   https://github.com/ipython/ipython/wiki/IPEP-16%3A-Notebook-multi-directory-dashboard-and-URL-mapping
+ * https://github.com/ipython/ipython/wiki/IPEP-16%3A-Notebook-multi-directory-dashboard-and-URL-mapping
  *
  * @author vlan
  */
@@ -33,7 +33,7 @@ public class WebSocketConnectionTest extends TestCase {
         assertTrue(connection.getKernelId().length() > 0);
         connection.shutdown();
       }
-    }, null);
+    }, null, DefaultProjectFactory.getInstance().getDefaultProject(), "");
     connection.close();
   }
 
@@ -60,7 +60,7 @@ public class WebSocketConnectionTest extends TestCase {
           connection.shutdown();
         }
       }
-    }, null);
+    }, null, DefaultProjectFactory.getInstance().getDefaultProject(), "");
     connection.close();
     assertTrue(evaluated.get());
   }
@@ -94,7 +94,7 @@ public class WebSocketConnectionTest extends TestCase {
           connection.shutdown();
         }
       }
-    }, null);
+    }, null, DefaultProjectFactory.getInstance().getDefaultProject(), "");
     connection.close();
     assertTrue(evaluated.get());
   }