1# -*- test-case-name: twisted.conch.test.test_ssh -*-
2# Copyright (c) Twisted Matrix Laboratories.
3# See LICENSE for details.
4
5
6"""
7Common functions for the SSH classes.
8
9Maintainer: Paul Swartz
10"""
11
12import struct, warnings, __builtin__
13
14try:
15    from Crypto import Util
16except ImportError:
17    warnings.warn("PyCrypto not installed, but continuing anyways!",
18            RuntimeWarning)
19
20from twisted.python import randbytes
21
22
23def NS(t):
24    """
25    net string
26    """
27    return struct.pack('!L',len(t)) + t
28
29def getNS(s, count=1):
30    """
31    get net string
32    """
33    ns = []
34    c = 0
35    for i in range(count):
36        l, = struct.unpack('!L',s[c:c+4])
37        ns.append(s[c+4:4+l+c])
38        c += 4 + l
39    return tuple(ns) + (s[c:],)
40
41def MP(number):
42    if number==0: return '\000'*4
43    assert number>0
44    bn = Util.number.long_to_bytes(number)
45    if ord(bn[0])&128:
46        bn = '\000' + bn
47    return struct.pack('>L',len(bn)) + bn
48
49def getMP(data, count=1):
50    """
51    Get multiple precision integer out of the string.  A multiple precision
52    integer is stored as a 4-byte length followed by length bytes of the
53    integer.  If count is specified, get count integers out of the string.
54    The return value is a tuple of count integers followed by the rest of
55    the data.
56    """
57    mp = []
58    c = 0
59    for i in range(count):
60        length, = struct.unpack('>L',data[c:c+4])
61        mp.append(Util.number.bytes_to_long(data[c+4:c+4+length]))
62        c += 4 + length
63    return tuple(mp) + (data[c:],)
64
65def _MPpow(x, y, z):
66    """return the MP version of (x**y)%z
67    """
68    return MP(pow(x,y,z))
69
70def ffs(c, s):
71    """
72    first from second
73    goes through the first list, looking for items in the second, returns the first one
74    """
75    for i in c:
76        if i in s: return i
77
78getMP_py = getMP
79MP_py = MP
80_MPpow_py = _MPpow
81pyPow = pow
82
83def _fastgetMP(data, count=1):
84    mp = []
85    c = 0
86    for i in range(count):
87        length = struct.unpack('!L', data[c:c+4])[0]
88        mp.append(long(gmpy.mpz(data[c + 4:c + 4 + length][::-1] + '\x00', 256)))
89        c += length + 4
90    return tuple(mp) + (data[c:],)
91
92def _fastMP(i):
93    i2 = gmpy.mpz(i).binary()[::-1]
94    return struct.pack('!L', len(i2)) + i2
95
96def _fastMPpow(x, y, z=None):
97    r = pyPow(gmpy.mpz(x),y,z).binary()[::-1]
98    return struct.pack('!L', len(r)) + r
99
100def install():
101    global getMP, MP, _MPpow
102    getMP = _fastgetMP
103    MP = _fastMP
104    _MPpow = _fastMPpow
105    # XXX: We override builtin pow so that PyCrypto can benefit from gmpy too.
106    def _fastpow(x, y, z=None, mpz=gmpy.mpz):
107        if type(x) in (long, int):
108            x = mpz(x)
109        return pyPow(x, y, z)
110    __builtin__.pow = _fastpow # evil evil
111
112try:
113    import gmpy
114    install()
115except ImportError:
116    pass
117
118