
import unittest

import const
import uniformdh

import obfsproxy.network.buffer as obfs_buf

class UniformDHTest( unittest.TestCase ):

    def setUp( self ):
        weAreServer = True
        self.udh = uniformdh.new("A" * const.SHARED_SECRET_LENGTH, weAreServer)

    def test1_createHandshake( self ):
        handshake = self.udh.createHandshake()
        self.failUnless((const.PUBLIC_KEY_LENGTH +
                         const.MARKER_LENGTH +
                         const.HMAC_LENGTH) <= len(handshake) <=
                        (const.MARKER_LENGTH +
                         const.HMAC_LENGTH +
                         const.MAX_PADDING_LENGTH))

    def test2_receivePublicKey( self ):
        buf = obfs_buf.Buffer(self.udh.createHandshake())

        self.failUnless(self.udh.receivePublicKey(buf, lambda x: x) == True)

        publicKey = self.udh.getRemotePublicKey()
        self.failUnless(len(publicKey) == const.PUBLIC_KEY_LENGTH)

    def test3_invalidHMAC( self ):
        # Make the HMAC invalid.
        handshake = self.udh.createHandshake()
        if handshake[-1] != 'a':
            handshake = handshake[:-1] + 'a'
        else:
            handshake = handshake[:-1] + 'b'

        buf = obfs_buf.Buffer(handshake)

        self.failIf(self.udh.receivePublicKey(buf, lambda x: x) == True)


if __name__ == '__main__':
    unittest.main()


import unittest

import const
import uniformdh

import obfsproxy.network.buffer as obfs_buf

class UniformDHTest( unittest.TestCase ):

    def setUp( self ):
        weAreServer = True
        self.udh = uniformdh.new("A" * const.SHARED_SECRET_LENGTH, weAreServer)

    def test1_createHandshake( self ):
        handshake = self.udh.createHandshake()
        self.failUnless((const.PUBLIC_KEY_LENGTH +
                         const.MARKER_LENGTH +
                         const.HMAC_LENGTH) <= len(handshake) <=
                        (const.MARKER_LENGTH +
                         const.HMAC_LENGTH +
                         const.MAX_PADDING_LENGTH))

    def test2_receivePublicKey( self ):
        buf = obfs_buf.Buffer(self.udh.createHandshake())

        def callback( masterKey ):
            self.failUnless(len(masterKey) == const.MASTER_KEY_LENGTH)

        self.failUnless(self.udh.receivePublicKey(buf, callback) == True)

        publicKey = self.udh.getRemotePublicKey()
        self.failUnless(len(publicKey) == const.PUBLIC_KEY_LENGTH)

    def test3_invalidHMAC( self ):
        # Make the HMAC invalid.
        handshake = self.udh.createHandshake()
        if handshake[-1] != 'a':
            handshake = handshake[:-1] + 'a'
        else:
            handshake = handshake[:-1] + 'b'

        buf = obfs_buf.Buffer(handshake)

        self.failIf(self.udh.receivePublicKey(buf, lambda x: x) == True)


if __name__ == '__main__':
    unittest.main()

Extend UniformDH test to also verify the length of the shared master secret.
Extend UniformDH test to also verify the length of the shared master secret.

  
  import unittest
  
  import const
  import uniformdh
  
  import obfsproxy.network.buffer as obfs_buf
  
  class UniformDHTest( unittest.TestCase ):
  
      def setUp( self ):
          weAreServer = True
          self.udh = uniformdh.new("A" * const.SHARED_SECRET_LENGTH, weAreServer)
  
      def test1_createHandshake( self ):
          handshake = self.udh.createHandshake()
          self.failUnless((const.PUBLIC_KEY_LENGTH +
                           const.MARKER_LENGTH +
                           const.HMAC_LENGTH) <= len(handshake) <=
                          (const.MARKER_LENGTH +
                           const.HMAC_LENGTH +
                           const.MAX_PADDING_LENGTH))
  
      def test2_receivePublicKey( self ):
          buf = obfs_buf.Buffer(self.udh.createHandshake())
  
+         def callback( masterKey ):
+             self.failUnless(len(masterKey) == const.MASTER_KEY_LENGTH)
+ 
-         self.failUnless(self.udh.receivePublicKey(buf, lambda x: x) == True)
+         self.failUnless(self.udh.receivePublicKey(buf, callback) == True)
  
          publicKey = self.udh.getRemotePublicKey()
          self.failUnless(len(publicKey) == const.PUBLIC_KEY_LENGTH)
  
      def test3_invalidHMAC( self ):
          # Make the HMAC invalid.
          handshake = self.udh.createHandshake()
          if handshake[-1] != 'a':
              handshake = handshake[:-1] + 'a'
          else:
              handshake = handshake[:-1] + 'b'
  
          buf = obfs_buf.Buffer(handshake)
  
          self.failIf(self.udh.receivePublicKey(buf, lambda x: x) == True)
  
  
  if __name__ == '__main__':
      unittest.main()
  
Extend UniformDH test to also verify the length of the shared master secret.
## Code Before:

import unittest

import const
import uniformdh

import obfsproxy.network.buffer as obfs_buf

class UniformDHTest( unittest.TestCase ):

    def setUp( self ):
        weAreServer = True
        self.udh = uniformdh.new("A" * const.SHARED_SECRET_LENGTH, weAreServer)

    def test1_createHandshake( self ):
        handshake = self.udh.createHandshake()
        self.failUnless((const.PUBLIC_KEY_LENGTH +
                         const.MARKER_LENGTH +
                         const.HMAC_LENGTH) <= len(handshake) <=
                        (const.MARKER_LENGTH +
                         const.HMAC_LENGTH +
                         const.MAX_PADDING_LENGTH))

    def test2_receivePublicKey( self ):
        buf = obfs_buf.Buffer(self.udh.createHandshake())

        self.failUnless(self.udh.receivePublicKey(buf, lambda x: x) == True)

        publicKey = self.udh.getRemotePublicKey()
        self.failUnless(len(publicKey) == const.PUBLIC_KEY_LENGTH)

    def test3_invalidHMAC( self ):
        # Make the HMAC invalid.
        handshake = self.udh.createHandshake()
        if handshake[-1] != 'a':
            handshake = handshake[:-1] + 'a'
        else:
            handshake = handshake[:-1] + 'b'

        buf = obfs_buf.Buffer(handshake)

        self.failIf(self.udh.receivePublicKey(buf, lambda x: x) == True)


if __name__ == '__main__':
    unittest.main()

## Instruction:
Extend UniformDH test to also verify the length of the shared master secret.
## Code After:

import unittest

import const
import uniformdh

import obfsproxy.network.buffer as obfs_buf

class UniformDHTest( unittest.TestCase ):

    def setUp( self ):
        weAreServer = True
        self.udh = uniformdh.new("A" * const.SHARED_SECRET_LENGTH, weAreServer)

    def test1_createHandshake( self ):
        handshake = self.udh.createHandshake()
        self.failUnless((const.PUBLIC_KEY_LENGTH +
                         const.MARKER_LENGTH +
                         const.HMAC_LENGTH) <= len(handshake) <=
                        (const.MARKER_LENGTH +
                         const.HMAC_LENGTH +
                         const.MAX_PADDING_LENGTH))

    def test2_receivePublicKey( self ):
        buf = obfs_buf.Buffer(self.udh.createHandshake())

        def callback( masterKey ):
            self.failUnless(len(masterKey) == const.MASTER_KEY_LENGTH)

        self.failUnless(self.udh.receivePublicKey(buf, callback) == True)

        publicKey = self.udh.getRemotePublicKey()
        self.failUnless(len(publicKey) == const.PUBLIC_KEY_LENGTH)

    def test3_invalidHMAC( self ):
        # Make the HMAC invalid.
        handshake = self.udh.createHandshake()
        if handshake[-1] != 'a':
            handshake = handshake[:-1] + 'a'
        else:
            handshake = handshake[:-1] + 'b'

        buf = obfs_buf.Buffer(handshake)

        self.failIf(self.udh.receivePublicKey(buf, lambda x: x) == True)


if __name__ == '__main__':
    unittest.main()
